Update train.py

train.py

@@ -1,26 +1,44 @@
 import torch
 import torch.nn as nn
-from pathlib import Path
-from data_prep import test_loader, device
+from data_prep import train_loader, val_loader, device
 from models.model import PlantCNN
 from utils.config import load_config
-from clearml import Task, InputModel
-import numpy as np
-from utils.vis import visualize_preds, plot_cfm
+from clearml import Task
+from pathlib import Path
 from tqdm.auto import tqdm
-import ast
 
+def train_step(model, loader, optimizer, loss_fn, device):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    
+    for batch_idx, batch in enumerate(tqdm(loader, desc="Train", leave=False)):
+        images = batch["pixel_values"].to(device)
+        labels = batch["labels"].to(device)
+        
+        optimizer.zero_grad()
+        output = model(images)
+        loss = loss_fn(output, labels)
+        loss.backward()
+        optimizer.step()
+        
+        running_loss += loss.item()*labels.size(0)
+        
+        _, preds = torch.max(output, dim=1)
+        correct += (preds==labels).sum().item()
+        total += labels.size(0)
+        
+    epoch_loss = running_loss/total
+    epoch_acc = correct/total
+    return epoch_loss, epoch_acc
 
-def evaluate_on_test(model, loader, loss_fn, device, num_imgs):
+def test_step(model, loader, loss_fn, device):
     model.eval()
-    all_labels = []
-    all_preds = []
+    
     running_loss = 0.0
     correct = 0
     total = 0
-    imgs_to_display = []
-    lbls_to_display = []
-    prs_to_display = []
     
     with torch.no_grad():
         for batch_idx, batch in enumerate(tqdm(loader, desc="Val", leave=False)):
@@ -36,69 +54,75 @@ def evaluate_on_test(model, loader, loss_fn, device, num_imgs):
             correct += (preds==labels).sum().item()
             total += labels.size(0)
             
-            all_labels.extend(labels.cpu().numpy())
-            all_preds.extend(preds.cpu().numpy())
-            
-            if len(imgs_to_display) < num_imgs:
-                remaining = num_imgs - len(imgs_to_display)
-                for img, lbl, pr in zip(images[:remaining], preds[:remaining], preds[:remaining]):
-                    imgs_to_display.append(img.cpu())
-                    lbls_to_display.append(lbl.item())
-                    prs_to_display.append(pr.item())
-                    
-    test_loss = running_loss / total
-    test_acc = correct / total
-    return test_loss, test_acc, all_labels, all_preds, imgs_to_display, lbls_to_display, prs_to_display
-    
-
-
+    epoch_loss = running_loss/total
+    epoch_acc = correct/total
+    return epoch_loss, epoch_acc
 
 def main():
+    config = load_config()
+    num_classes = config["num_classes"]
+    channels = config["channels"]
+    dropout = config["dropout"]
+    lr = config["lr"]
+    weight_decay = config["weight_decay"]
+    num_epochs = config["num_epochs"]
+    patience = config["early_stopping_patience"]
     project_name = "GAP_plant_disease_classification"
-    model_name = "PlantCNN"
+    model_name="PlantCNN"
     
-    task = Task.init(project_name=project_name, task_name=f"{model_name}_test")
+    model = PlantCNN(num_classes=num_classes, channels=channels, dropout=dropout).to(device)
+    loss_fn = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+
+    task = Task.init(project_name=project_name, task_name=f"{model_name}_training")
+    task.connect(config)
+    task.add_tags([model_name, "train"])
     logger = task.get_logger()
     
-    input_model = InputModel(model_id="b9308022b85e4eea952d78124d1ee597")
-    
-    training_task_id = input_model.task
-    training_task = Task.get_task(task_id=training_task_id)
-    
-    training_params = training_task.get_parameters()
-    print(f"Training parameters: {training_params}")
+    best_val_acc = 0.0
+    best_state_dict = None
+    patience_cnt = 0
     
-    num_classes = int(training_params.get("General/num_classes"))
-    channels = ast.literal_eval(training_params.get("General/channels"))
-    dropout = float(training_params.get("General/dropout"))
-    mean_nm = ast.literal_eval(training_params.get("General/normalize_mean"))
-    std_nm = ast.literal_eval(training_params.get("General/normalize_std"))
-    kernel_sizes = ast.literal_eval(training_params.get("General/kernel_sizes"))
+    for epoch in range(num_epochs):
+        print(f"\nEpoch: {epoch+1}/{num_epochs}")
+        
+        train_loss, train_acc = train_step(
+            model, train_loader, optimizer, loss_fn, device
+        )
+        val_loss, val_acc = test_step(
+            model, val_loader, loss_fn, device
+        )
     
-    task.add_tags([model_name, "test"])
-    
-    dataset = test_loader.dataset
-    class_names = dataset.features["label"].names
-
-
-    model = PlantCNN(num_classes=num_classes, channels=channels, dropout=dropout, kernel_sizes=kernel_sizes).to(device)
-    model_path = input_model.get_local_copy()
-    state_dict = torch.load(model_path, map_location=device)
-    model.load_state_dict(state_dict)
-    
-    loss_fn = nn.CrossEntropyLoss()
+        print(f"Train loss: {train_loss:.3f} | Train accuracy: {train_acc:.3f}")
+        print(f"Validation loss: {val_loss:.3f} | Validation accuracy: {val_acc:.3f}")
+        logger.report_scalar("loss", "train", train_loss, epoch)
+        logger.report_scalar("loss", "val", val_loss, epoch)
+        logger.report_scalar("accuracy", "train", train_acc, epoch)
+        logger.report_scalar("accuracy", "val", val_acc, epoch)
     
-    test_loss, test_acc, all_labels, all_preds, display_images, display_labels, display_preds = evaluate_on_test(model, test_loader,
-                                                                                                                 loss_fn, device,
-                                                                                                                 num_imgs=24)
+        if val_acc > best_val_acc:
+            best_val_acc = val_acc
+            best_state_dict = model.state_dict()
+            patience_cnt = 0
+        else:
+            patience_cnt+=1
+        
+        if patience_cnt >= patience:
+            print(f"\nEarly stopping triggered after {epoch+1} epochs.")
+            break
+        
+    if best_state_dict is not None:
+        model.load_state_dict(best_state_dict)
+        
+    project_rt = Path(__file__).resolve().parent
+    model_dir = project_rt/"saved_models"
+    model_dir.mkdir(parents=True, exist_ok=True)
+    model_path = model_dir/"plant_cnn.pt"
     
-    print("\nTest results:")
-    print(f"Test loss: {test_loss:.3f} | Test accuracy: {test_acc:.3f}")
-    logger.report_scalar("loss", "test", test_loss, 0)
-    logger.report_scalar("accuracy", "test", test_acc, 0)
+    torch.save(model.state_dict(), model_path)
+    print(f"Saved best model to {model_path}")
     
-    visualize_preds(display_images, display_labels, display_preds, logger, class_names, mean_nm, std_nm, num_images=24)
-    plot_cfm(all_labels, all_preds, logger, class_names, num_classes)
+    task.update_output_model(model_path=str(model_path), name="plant_cnn_best")
     
 if __name__ == "__main__":
-    main()
+    main()