update docs and trainer

README.md

@@ -157,16 +157,16 @@ from data4co.draw.tsp import draw_tsp_solution, draw_tsp_problem
 
 # use TSPConcordeSolver to solve the problem
 solver = TSPConcordeSolver(scale=1)
-solver.from_tsp("docs/kroA150.tsp")
+solver.from_tsp("examples/tsp/kroA150.tsp")
 solver.solve(norm="EUC_2D")
 
 # draw
 draw_tsp_problem(
-    save_path="docs/kroA150_problem.png",
+    save_path="docs/assets/kroA150_problem.png",
     points=solver.ori_points,
 )
 draw_tsp_solution(
-    save_path="docs/kroA150_solution.png",
+    save_path="docs/assets/kroA150_solution.png",
     points=solver.ori_points,
     tours=solver.tours
 )
@@ -187,17 +187,17 @@ from data4co import draw_mis_problem, draw_mis_solution
 
 # use KaMISSolver to solve the problem
 mis_solver = KaMISSolver()
-mis_solver.solve(src="docs/mis_example")
+mis_solver.solve(src="examples/mis_example")
 
 # draw
 draw_mis_problem(
-    save_path="docs/mis_problem.png", 
-    ckle_path="docs/mis_example/mis_example.gpickle"
+    save_path="docs/assets/mis_problem.png", 
+    gpickle_path="examples/mis/mis_example.gpickle"
 )
 draw_mis_solution(
     save_path="docs/mis_solution.png",
-    gpickle_path="docs/mis_example/mis_example.gpickle",
-    result_path="docs/mis_example/solve/mis_example_unweighted.result"
+    gpickle_path="examples/mis/mis_example.gpickle",
+    result_path="examples/mis/solve/mis_example_unweighted.result"
 )
 ```
 

docs/project_example/data.py

@@ -1,9 +1,5 @@
-import os
 import torch
-import pickle
 import numpy as np
-from sklearn.neighbors import KDTree
-from torch_geometric.data import Data as GraphData
 
 
 class TSPDataset(torch.utils.data.Dataset):

docs/project_example/env.py

@@ -1,7 +1,7 @@
 import torch
 import torch.utils.data
-from ml4co_kit.learning.env import BaseEnv
 from data import TSPDataset
+from ml4co_kit.learning.env import BaseEnv
 from torch_geometric.data import DataLoader as GraphDataLoader
 
 

docs/project_example/model.py

@@ -1,12 +1,12 @@
 import torch
 import numpy as np
-from typing import Any, Union
+import torch.nn.functional as F
 from ml4co_kit.learning.model import BaseEncoder
 from ml4co_kit.evaluate import TSPEvaluator
-from env import TSPEnv
-import torch.nn.functional as F
 from ml4co_kit.learning.search import SearchConfigurator
+from env import TSPEnv
 from search import tsp_greedy, tsp_2opt
+from typing import Any, Union
 
 
 class TSPGNN(BaseEncoder):

docs/project_example/search.py

@@ -1,7 +1,13 @@
 import numpy as np
 
 
-def tsp_greedy(adj_mat, np_points, parallel_sampling=1, device="cpu", **kwargs):
+def tsp_greedy(
+    adj_mat: np.ndarray, 
+    np_points: np.ndarray, 
+    parallel_sampling: int=1, 
+    device: str="cpu", 
+    **kwargs
+):
     raise NotImplementedError
 
 

docs/project_example/solver.py

@@ -1,5 +1,4 @@
 import time
-import torch
 import numpy as np
 from tqdm import tqdm
 from typing import Union, Any
@@ -84,6 +83,12 @@ def solve(
         decoded_tours = list()
         for idx in tqdm(range(self.points.shape[0]), desc="Decoding"):
             adj_mat = np.expand_dims(heatmap[idx], axis=0)
+            tour = self.decoding_func(
+                adj_mat=adj_mat, 
+                np_points=self.points[idx], 
+                edge_index_np=None,
+                **decoding_kwargs
+            )
             decoded_tours.append(tour[0])
         decoded_tours = np.array(decoded_tours)
 

ml4co_kit/learning/trainer.py

@@ -1,16 +1,14 @@
 import os
-from wandb.util import generate_id
-from typing import Optional
 import torch
 from torch import nn
+from typing import Optional
+from wandb.util import generate_id
 from typing import Union, Optional
-from pytorch_lightning import Trainer
-from pytorch_lightning.strategies.strategy import Strategy
-from pytorch_lightning.strategies.ddp import DDPStrategy
-from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
-from pytorch_lightning.callbacks.progress.tqdm_progress import TQDMProgressBar
+from pytorch_lightning.trainer import Trainer
+from pytorch_lightning.strategies import Strategy, DDPStrategy
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint, TQDMProgressBar
 from pytorch_lightning.utilities import rank_zero_info
-from pytorch_lightning.loggers.wandb import WandbLogger
+from pytorch_lightning.loggers import WandbLogger
 
 
 class Checkpoint(ModelCheckpoint):
@@ -46,7 +44,7 @@ def __init__(
         resume_id: Optional[str] = None,
     ):
         if not os.path.exists(save_dir):
-            os.mkdir(save_dir)
+            os.makedirs(save_dir)
         if id is None and resume_id is None:
             wandb_id = os.getenv("WANDB_RUN_ID") or generate_id()
         else:

setup.py

@@ -96,7 +96,7 @@ def get_tag(self):
     author=AUTHOR,
     url=URL,
     packages=find_packages(exclude=["tests", "*.tests", "*.tests.*", "tests.*"]),
-    package_data={NAME: ["**"]},
+    package_data={NAME: ["**"], "docs": ["**"]},
     install_requires=REQUIRED,
     extras_require=EXTRAS,
     include_package_data=True,