update

.gitignore

@@ -163,11 +163,14 @@ cython_debug/
 # ignore dataset but not the folder
 data/raw/*
 data/processed/*
+results/models/*
 
 !data/raw/.gitkeep
 !data/processed/.gitkeep
+!results/models/.gitkeep
 
-uvenv/
+# tensorboard logs
 runs/
-wandb/
-ubuntu-venv/
+
+# wandb logs
+wandb/

Makefile

@@ -1,4 +1,6 @@
 package:
 	pip freeze > requirements.txt
 venv:
-	source /mnt/d/ubuntu/env/mlenv/bin/activate
+	source /mnt/d/ubuntu/env/mlenv/bin/activate
+tensorboard:
+	tensorboard --inspect --logdir logs/tensorboard

app.py

@@ -18,7 +18,7 @@ def classify_drawing(drawing_image):
     num_classes = 3  # Set the number of classes
     # Initialize your model class
     model = ShapeClassifier(num_classes=num_classes)
-    model.load_state_dict(torch.load('model.pth', map_location=torch.device('cpu')))
+    model.load_state_dict(torch.load('results/models/model.pth', map_location=torch.device('cpu')))
     model.eval()  # Set the model to evaluation mode
 
     # Convert the drawing to a grayscale image
@@ -45,6 +45,5 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),  # Use Sketchpad as input
     outputs="text",
     live=True,
-    capture_session=True,
 )
 iface.launch(server_port=7860)

main.py

@@ -1,10 +1,47 @@
-from src.train import train_runner
-from src.auto import auto_hyper_parameter
-import os
-# set WANDB_API_KEY=$YOUR_API_KEY
-# os.environ["WANDB_API_KEY"] = '7c0f2b9470a0a5c82bfae5bab4705344cb53288b'
-# os.environ['WANDB_MODE'] = "offline"
+
+
+import torch
+import torch.optim as optim
+import torch.nn.functional as F
+from src.models.model import ShapeClassifier
+
+from src.configs.model_config import ModelConfig
+from src.data.data_loader import train_loader, num_classes
+from src.utils.tensorboard import writer
+from src.utils.train import train
+from src.utils.test import test
+from src.utils.wandb import wandb
+from src.utils.logs import logging
+import json
+from src.utils.model import save_model
+
+
+def train_runner():
+
+    config = ModelConfig().get_config()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = ShapeClassifier(num_classes=num_classes).to(device)
+    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)
+    # log models config to wandb
+    wandb.config.update(config)
+
+    for epoch in range(config.epochs):
+        print(f"Epoch {epoch+1}\n-------------------------------")
+
+        loss = train(train_loader, model=model, loss_fn=F.cross_entropy,
+                     optimizer=optimizer)
+        test(train_loader, model=model, loss_fn=F.cross_entropy)
+        # 3. Log metrics over time to visualize performance
+        wandb.log({"loss": loss})
+
+        # save model
+        save_model(model, "results/models/last.pth")
+
+        # 4. Log an artifact to W&B
+        wandb.log_artifact("model.pth")
+        # model.train()
+
+
 if __name__ == "__main__":
-    print("Training the model...")
-    # train_runner()
-    auto_hyper_parameter()
+    logging.info("Training model")
+    train_runner()

src/configs/model_config.py

@@ -1,11 +1,19 @@
 import torch
 
+
 class ModelConfig:
     def __init__(self):
         self.learning_rate = 0.001
         self.batch_size = 32
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.epochs = 5
-        self.log_interval = 2 # Log every 2 batches => number of items is 32*2 = 64
+        self.log_interval = 2  # Log every 2 batches => number of items is 32*2 = 64
+        
+        # Wandb config
+        self.wandb = True
+        self.wandb_project = "template-pytorch-model"
+        self.wandb_entity = "nguyen"
+        self.wandb_api_key = ""
+
     def get_config(self):
-        return self
+        return self

src/train.py

@@ -1,40 +0,0 @@
-import torch
-import torch.optim as optim
-import torch.nn.functional as F
-from src.models.model import ShapeClassifier
-
-from src.configs.model_config import ModelConfig
-from src.data.data_loader import train_loader, num_classes
-from src.utils.logs import writer
-from src.utils.train import train
-from src.utils.test import test
-import wandb
-import json
-wandb.init(project="template-pytorch-model", entity="nguyen")
-
-
-def train_runner():
-
-    config = ModelConfig().get_config()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = ShapeClassifier(num_classes=num_classes).to(device)
-    optimizer = optim.Adam(model.parameters(), lr=config.learning_rate)
-    log_interval = 20
-    # log models config to wandb
-    wandb.config.update(config)
-
-    for epoch in range(config.epochs):
-        print(f"Epoch {epoch+1}\n-------------------------------")
-
-        loss = train(train_loader, model=model, loss_fn=F.cross_entropy,
-                     optimizer=optimizer)
-        test(train_loader, model=model, loss_fn=F.cross_entropy)
-        # 3. Log metrics over time to visualize performance
-        wandb.log({"loss": loss})
-
-        # save model
-        torch.save(model.state_dict(), "model.pth")
-
-        # 4. Log an artifact to W&B
-        wandb.log_artifact("model.pth")
-        # model.train()

src/utils/__init__.py

@@ -1,17 +1,8 @@
-import importlib
-import os
-from inspect import isclass
 
-# import all files under utils/
-utils_dir = os.path.dirname(__file__)
-for file in os.listdir(utils_dir):
-    path = os.path.join(utils_dir, file)
-    if not file.startswith("_") and not file.startswith(".") and (file.endswith(".py") or os.path.isdir(path)):
-        config_name = file[: file.find(".py")] if file.endswith(".py") else file
-        module = importlib.import_module("src.utils." + config_name)
-        for attribute_name in dir(module):
-            attribute = getattr(module, attribute_name)
 
-            if isclass(attribute):
-                # Add the class to this package's variables
-                globals()[attribute_name] = attribute
+# import all files under utils/
+from os.path import dirname, basename, isfile, join
+import glob
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [basename(f)[:-3] for f in modules if isfile(f)
+           and not f.endswith('__init__.py')]

src/utils/file.py

@@ -0,0 +1,8 @@
+
+import os
+# enssure folder exists
+
+
+def ensure_folder_exists(path):
+    if not os.path.exists(path):
+        os.makedirs(path)

src/utils/logs.py

@@ -1,2 +1,34 @@
-from torch.utils.tensorboard import SummaryWriter
-writer = SummaryWriter()
+import os
+import logging
+
+log_dir = "logs"
+log_level = logging.INFO
+
+if not os.path.exists(log_dir):
+    os.makedirs(log_dir)
+
+log_filename = os.path.join(log_dir, "app.log")
+
+logging.basicConfig(
+    filename=log_filename,
+    level=log_level,
+    format="%(asctime)s [%(levelname)s]: %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S"
+)
+
+console_handler = logging.StreamHandler()
+console_handler.setLevel(log_level)
+console_handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s]: %(message)s"))
+logging.getLogger().addHandler(console_handler)
+
+def log_info(message):
+    logging.info(message)
+
+def log_warning(message):
+    logging.warning(message)
+
+def log_error(message):
+    logging.error(message)
+
+def log_exception(message):
+    logging.exception(message)

src/utils/model.py

@@ -0,0 +1,15 @@
+# save model utils
+import torch
+import torch
+import os
+
+
+def save_model(model: torch.nn.Module, path: str) -> str:
+    parent_folder = os.path.dirname(path)
+    os.makedirs(parent_folder, exist_ok=True)
+    torch.save(model.state_dict(), path)
+    return path
+
+def load_model(model: torch.nn.Module, path: str) -> torch.nn.Module:
+    model.load_state_dict(torch.load(path))
+    return model

src/utils/tensorboard.py

@@ -0,0 +1,2 @@
+from torch.utils.tensorboard import SummaryWriter
+writer = SummaryWriter()

src/utils/wandb.py

@@ -0,0 +1,17 @@
+import wandb
+from .logs import log_info
+from src.configs.model_config import ModelConfig
+config = ModelConfig().get_config()
+
+if config.wandb:
+    project = config.wandb_project
+    entity = config.wandb_entity
+    api_key = config.wandb_api_key
+    wandb.login(key=api_key)
+    wandb.init(project=project, entity=entity)
+    log_info("Wandb is enabled")
+else:
+    log_info("Wandb is disabled")
+    # disable wandb
+    wandb.init(mode="disabled")
+    pass