Modular code and removed misclassified images collection

main.py

@@ -5,11 +5,15 @@ from resnet_model import ResNet50
 from data_utils import get_train_transform, get_test_transform, get_data_loaders
 from train_test import train, test
 from utils import save_checkpoint, load_checkpoint, plot_training_curves, plot_misclassified_samples
+from torchsummary import summary
 
 def main():
     # Initialize model, loss function, and optimizer
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = ResNet50().to(device)
+    model = ResNet50()
+    model = torch.nn.DataParallel(model)
+    model = model.to(device)
+    summary(model, input_size=(3, 224, 224))
     criterion = nn.CrossEntropyLoss()
     optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
 
@@ -56,7 +60,9 @@ def main():
     plot_training_curves(epochs, train_acc1, test_acc1, train_acc5, test_acc5, train_losses, test_losses, learning_rates)
 
     # Plot misclassified samples
+    '''
     plot_misclassified_samples(misclassified_images, misclassified_labels, misclassified_preds, classes=['class1', 'class2', ...])  # Replace with actual class names
+    '''
 
 if __name__ == '__main__':
-    main() 
+    main() 

tmppl87qjev/_remote_module_non_scriptable.py

@@ -1,81 +0,0 @@
-from typing import *
-
-import torch
-import torch.distributed.rpc as rpc
-from torch import Tensor
-from torch._jit_internal import Future
-from torch.distributed.rpc import RRef
-from typing import Tuple  # pyre-ignore: unused import
-
-
-module_interface_cls = None
-
-
-def forward_async(self, *args, **kwargs):
-    args = (self.module_rref, self.device, self.is_device_map_set, *args)
-    kwargs = {**kwargs}
-    return rpc.rpc_async(
-        self.module_rref.owner(),
-        _remote_forward,
-        args,
-        kwargs,
-    )
-
-
-def forward(self, *args, **kwargs):
-    args = (self.module_rref, self.device, self.is_device_map_set, *args)
-    kwargs = {**kwargs}
-    ret_fut = rpc.rpc_async(
-        self.module_rref.owner(),
-        _remote_forward,
-        args,
-        kwargs,
-    )
-    return ret_fut.wait()
-
-
-_generated_methods = [
-    forward_async,
-    forward,
-]
-
-
-
-
-def _remote_forward(
-    module_rref: RRef[module_interface_cls], device: str, is_device_map_set: bool, *args, **kwargs):
-    module = module_rref.local_value()
-    device = torch.device(device)
-
-    if device.type != "cuda":
-        return module.forward(*args, **kwargs)
-
-    # If the module is on a cuda device,
-    # move any CPU tensor in args or kwargs to the same cuda device.
-    # Since torch script does not support generator expression,
-    # have to use concatenation instead of
-    # ``tuple(i.to(device) if isinstance(i, Tensor) else i for i in *args)``.
-    args = (*args,)
-    out_args: Tuple[()] = ()
-    for arg in args:
-        arg = (arg.to(device),) if isinstance(arg, Tensor) else (arg,)
-        out_args = out_args + arg
-
-    kwargs = {**kwargs}
-    for k, v in kwargs.items():
-        if isinstance(v, Tensor):
-            kwargs[k] = kwargs[k].to(device)
-
-    if is_device_map_set:
-        return module.forward(*out_args, **kwargs)
-
-    # If the device map is empty, then only CPU tensors are allowed to send over wire,
-    # so have to move any GPU tensor to CPU in the output.
-    # Since torch script does not support generator expression,
-    # have to use concatenation instead of
-    # ``tuple(i.cpu() if isinstance(i, Tensor) else i for i in module.forward(*out_args, **kwargs))``.
-    ret: Tuple[()] = ()
-    for i in module.forward(*out_args, **kwargs):
-        i = (i.cpu(),) if isinstance(i, Tensor) else (i,)
-        ret = ret + i
-    return ret

train_test.py

@@ -31,6 +31,9 @@ def train(model, device, train_loader, optimizer, criterion, epoch, accumulation
 
         pbar.set_description(desc=f'Epoch {epoch} | Loss: {running_loss / (batch_idx + 1):.4f} | Top-1 Acc: {100. * correct1 / total:.2f} | Top-5 Acc: {100. * correct5 / total:.2f}')
 
+        if (batch_idx + 1) % 50 == 0:
+            torch.cuda.empty_cache()
+
     return 100. * correct1 / total, 100. * correct5 / total, running_loss / len(train_loader)
 
 def test(model, device, test_loader, criterion):
@@ -56,13 +59,15 @@ def test(model, device, test_loader, criterion):
             correct5 += predicted.eq(targets.view(-1, 1).expand_as(predicted)).sum().item()
 
             # Collect misclassified samples
+            '''
             for i in range(inputs.size(0)):
                 if targets[i] not in predicted[i, :1]:
                     misclassified_images.append(inputs[i].cpu())
                     misclassified_labels.append(targets[i].cpu())
                     misclassified_preds.append(predicted[i, :1].cpu())
+            '''
 
     test_accuracy1 = 100. * correct1 / total
     test_accuracy5 = 100. * correct5 / total
     print(f'Test Loss: {test_loss/len(test_loader):.4f}, Top-1 Accuracy: {test_accuracy1:.2f}, Top-5 Accuracy: {test_accuracy5:.2f}')
-    return test_accuracy1, test_accuracy5, test_loss / len(test_loader), misclassified_images, misclassified_labels, misclassified_preds 
+    return test_accuracy1, test_accuracy5, test_loss / len(test_loader), misclassified_images, misclassified_labels, misclassified_preds 

utils.py

@@ -9,13 +9,15 @@ def save_checkpoint(model, optimizer, epoch, loss, path):
         'optimizer_state_dict': optimizer.state_dict(),
         'loss': loss,
     }, path)
+    print(f"Checkpoint saved at epoch {epoch}")
 
 def load_checkpoint(model, optimizer, path):
-    checkpoint = torch.load(path)
+    checkpoint = torch.load(path, weights_only=True)
     model.load_state_dict(checkpoint['model_state_dict'])
     optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
     epoch = checkpoint['epoch']
     loss = checkpoint['loss']
+    print(f"Checkpoint loaded, resuming from epoch {epoch}")
     return model, optimizer, epoch, loss
 
 def plot_training_curves(epochs, train_acc1, test_acc1, train_acc5, test_acc5, train_losses, test_losses, learning_rates):
@@ -62,4 +64,4 @@ def plot_misclassified_samples(misclassified_images, misclassified_labels, miscl
         plt.imshow(misclassified_grid.permute(1, 2, 0))
         plt.title("Misclassified Samples")
         plt.axis('off')
-        plt.show() 
+        plt.show()