prepared training of dem model from cli

Author: KennethTM

dem_autoencoder.ipynb

@@ -10,35 +10,49 @@
 pl.seed_everything(9999)
 random.seed(9999)
 
-#Load DEM data
-data = np.load("data/data.npz")
-dem_data = dem_scale(data["dem"])
+#Define main function
+def dem_training():
 
-train = dem_data[:8000]
-valid = dem_data[8000:]
+    #Load DEM data
+    data = np.load("data/data.npz")
+    dem_data = dem_scale(data["dem"])
 
-#Load mask data
-random_mask = np.load("data/lakes_random.npz")
-mask_data = random_mask["mask"]
+    train = dem_data[:8000]
+    valid = dem_data[8000:]
 
-random_mask_idx = random.choices(range(mask_data.shape[0]), k=valid.shape[0])
-valid_mask = mask_data[random_mask_idx]
+    #Load mask data
+    random_mask = np.load("data/lakes_random.npz")
+    mask_data = random_mask["mask"]
 
-#Create datasets
-train_dataset = DEMTrain(train, mask_data)
-valid_dataset = DEMValid(valid, valid_mask)
+    random_mask_idx = random.choices(range(mask_data.shape[0]), k=valid.shape[0])
+    valid_mask = mask_data[random_mask_idx]
 
-#Create dataloaders
-train_loader = DataLoader(train_dataset, batch_size=32, num_workers=8, shuffle = True)
-val_loader = DataLoader(valid_dataset, batch_size=32, num_workers=8, shuffle = False)
+    #Create datasets
+    train_dataset = DEMTrain(train, mask_data)
+    valid_dataset = DEMValid(valid, valid_mask)
 
-#Initiate model
-lake_model = AutoEncoder(init_features=8)
+    #Create dataloaders
+    train_loader = DataLoader(train_dataset, batch_size=32, num_workers=0, shuffle = True)
+    val_loader = DataLoader(valid_dataset, batch_size=32, num_workers=0, shuffle = False)
 
-#Initiate trainer
-checkpoint_callback = ModelCheckpoint(monitor="val_loss", save_last=True)
+    #Initiate models of differing complexity
+    dem_model_8 = AutoEncoder(init_features=8)
+    dem_model_16 = AutoEncoder(init_features=16)
+    dem_model_32 = AutoEncoder(init_features=32)
+    dem_model_64 = AutoEncoder(init_features=64)
 
-trainer = pl.Trainer(gpus=1, max_epochs=1000, callbacks=checkpoint_callback, precision=16)
+    #Initiate callbacks
+    checkpoint_callback = ModelCheckpoint(monitor="val_loss", save_last=True)
+
+    #Initiate trainer
+    trainer = pl.Trainer(gpus=1, max_epochs=1000, callbacks=checkpoint_callback)
+
+    #Train model
+    trainer.fit(dem_model_8, train_loader, val_loader)
+    trainer.fit(dem_model_16, train_loader, val_loader)
+    trainer.fit(dem_model_32, train_loader, val_loader)
+    trainer.fit(dem_model_64, train_loader, val_loader)
+
+if __name__ == "__main__":
+    dem_training()
 
-#Train model
-trainer.fit(lake_model, train_loader, val_loader)

dem_model_train.py

@@ -2,7 +2,6 @@
 #https://github.com/mateuszbuda/brain-segmentation-pytorch
 
 from collections import OrderedDict
-
 import torch
 import torch.nn as nn
 
@@ -72,23 +71,19 @@ def _block(in_channels, features, name):
                             in_channels=in_channels,
                             out_channels=features,
                             kernel_size=3,
-                            padding=1,
-                            #bias=False,
+                            padding=1
                         ),
                     ),
-                    #(name + "norm1", nn.BatchNorm2d(num_features=features)),
                     (name + "act1", nn.LeakyReLU()),
                     (
                         name + "conv2",
                         nn.Conv2d(
                             in_channels=features,
                             out_channels=features,
                             kernel_size=3,
-                            padding=1,
-                            #bias=False,
+                            padding=1
                         ),
                     ),
-                    #(name + "norm2", nn.BatchNorm2d(num_features=features)),
                     (name + "act2", nn.LeakyReLU()),
                 ]
             )