changed model used by app to have less compute

app.py

@@ -204,7 +204,7 @@ def create_app():
     '''
     # Used to serve with panel serve in command line
     save_dir = FILE_PATH + '/saved_models'
-    base_name = 'tiny_vgg'
+    base_name = 'tiny_vgg_less_compute'
 
     mod_path = f'{save_dir}/{base_name}_model.pth' # Path to the saved model state dict
     settings_path = f'{save_dir}/{base_name}_settings.yaml' # Path to the saved model kwargs

model_training/args.txt

@@ -1,12 +1,12 @@
 --num-workers
 0
 --num-epochs
-100
+300
 --batch-size
 100
 --learning-rate
 0.001
 --patience
-20
+50
 --min-delta
-0.0005
+0.0001

model_training/data_setup.py

@@ -47,17 +47,19 @@ def get_dataloaders(root: str,
         num_workers (int): Number of workers to use for multiprocessing. Default is 0.
     '''
 
-    # Get training and testing MNIST data
+   # Get training and testing MNIST data
     mnist_train = datasets.MNIST(root, download = True, train = True, 
-                                transform = TRAIN_TRANSFORMS)
+                                 transform = TRAIN_TRANSFORMS)
     mnist_test = datasets.MNIST(root, download = True, train = False, 
                                 transform = BASE_TRANSFORMS)
 
     # Create dataloaders
     if num_workers > 0:
         mp_context = utils.MP_CONTEXT
+        persistent_workers = True
     else:
         mp_context = None
+        persistent_workers = False
 
     train_dl = DataLoader(
         dataset = mnist_train,
@@ -65,7 +67,8 @@ def get_dataloaders(root: str,
         shuffle = True,
         num_workers = num_workers,
         multiprocessing_context = mp_context,
-        pin_memory = True
+        pin_memory = utils.PIN_MEM,
+        persistent_workers = persistent_workers
     )
 
     test_dl = DataLoader(
@@ -74,7 +77,8 @@ def get_dataloaders(root: str,
         shuffle = False,
         num_workers = num_workers,
         multiprocessing_context = mp_context,
-        pin_memory = True
+        pin_memory = utils.PIN_MEM,
+        persistent_workers = persistent_workers
     )
 
     return train_dl, test_dl

model_training/run_training.py

@@ -53,7 +53,7 @@ if __name__ == '__main__':
     # Set up saving directory and file name
     save_dir = '../saved_models'
 
-    base_name = 'tiny_vgg'
+    base_name = 'tiny_vgg_less_compute'
     mod_name = f'{base_name}_model.pth'
 
     # Get TinyVGG model
@@ -61,12 +61,13 @@ if __name__ == '__main__':
         'num_blks': 2,
         'num_convs': 2,
         'in_channels': 1,
-        'hidden_channels': 10,
-        'fc_hidden_dim': 64,
+        'hidden_channels': 5,
+        'fc_hidden_dim': 128,
         'num_classes': len(train_dl.dataset.classes)
     }
 
     vgg_mod = model.TinyVGG(**mod_kwargs).to(utils.DEVICE)
+    torch.compile(vgg_mod)
 
     # Save model kwargs and train settings
     with open(f'{save_dir}/{base_name}_settings.yaml', 'w') as f:

model_training/utils.py

@@ -10,12 +10,15 @@ import os
 if torch.cuda.is_available():
     DEVICE = torch.device('cuda')
     MP_CONTEXT = None
+    PIN_MEM = True
 elif torch.backends.mps.is_available():
     DEVICE = torch.device('mps')
     MP_CONTEXT = 'forkserver'
+    PIN_MEM = False
 else:
     DEVICE = torch.device('cpu')
     MP_CONTEXT = None
+    PIN_MEM = False
 
 
 #####################################
@@ -37,6 +40,7 @@ def set_seed(seed: int = 0):
     torch.cuda.manual_seed_all(seed)
     
     torch.use_deterministic_algorithms(True)
+    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
 
 def save_model(model: torch.nn.Module,
                save_dir: str, 

saved_models/tiny_vgg_less_compute_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72c2bf04c913dd30f77bb7dde8e4f9bd253533dbd7349475dbc2d0775b875a5e
+size 145606

saved_models/tiny_vgg_less_compute_settings.yaml

@@ -0,0 +1,14 @@
+mod_kwargs:
+  fc_hidden_dim: 128
+  hidden_channels: 5
+  in_channels: 1
+  num_blks: 2
+  num_classes: 10
+  num_convs: 2
+train_kwargs:
+  batch_size: 100
+  learning_rate: 0.001
+  min_delta: 0.0001
+  num_epochs: 300
+  num_workers: 0
+  patience: 50