cluster train py

train.py

@@ -2,55 +2,63 @@ import torch
 from data import LodopabDataset
 from trainer import Trainer
 from model import UNet
+from torchinfo import summary
+import functools
+from datetime import date
+from pathlib import Path
 
+
+# data locations
 cache_files = {
     'train':
-        '//hpc.isc.pad.ptb.de/hipcl/daten841/kaspar01/lodopab-ct/fbp/cache_lodopab_train_fbp.npy',
+        '/gpfs001/scratch/kaspar01/lodopab-ct/fbp/cache_lodopab_train_fbp.npy',
     'validation':
-        '//hpc.isc.pad.ptb.de/hipcl/daten841/kaspar01/lodopab-ct/fbp/cache_lodopab_validation_fbp.npy',
+        '/gpfs001/scratch/kaspar01/lodopab-ct/fbp/cache_lodopab_validation_fbp.npy',
     'test':
-        '//hpc.isc.pad.ptb.de/hipcl/daten841/kaspar01/lodopab-ct/fbp/cache_lodopab_test_fbp.npy'}
-
-ground_truth_data_loc = '//hpc.isc.pad.ptb.de/hipcl/daten841/kaspar01/lodopab-ct/'
+        '/gpfs001/scratch/kaspar01/lodopab-ct/fbp/cache_lodopab_test_fbp.npy'}
 
-# # cluster
-# cache_files = {
-#     'train':
-#         '/oneFS/daten841/kaspar01/lodopab-ct/fbp/cache_lodopab_train_fbp.npy',
-#     'validation':
-#         '/oneFS/daten841/kaspar01/lodopab-ct/fbp/cache_lodopab_validation_fbp.npy',
-#     'test':
-#         '/oneFS/daten841/kaspar01/lodopab-ct/fbp/cache_lodopab_test_fbp.npy'}
-#
-# ground_truth_data_loc = '/gpfs001/scratch/kaspar01/lodopab-ct/data/'
+ground_truth_data_loc = '/gpfs001/scratch/kaspar01/lodopab-ct/data'
 
 
-# auto define the correct device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Using {device} device.")
+model_path = Path('./model') / f'{date.today().strftime("%Y-%m-%d")}_unet.pt'
 
-training_dataset = LodopabDataset(cache_files=cache_files,
+# create Dataloaders
+training_dataset = LodopabDataset(cache_files=small_cache_files,
                                   ground_truth_data_loc=ground_truth_data_loc,
-                                  split='train')
+                                  split='train',
+                                  transform=functools.partial(torch.unsqueeze, dim=0),
+                                  target_transform=functools.partial(torch.unsqueeze, dim=0))
 
 training_dataloader = torch.utils.data.DataLoader(dataset=training_dataset,
                                                   batch_size=16,
                                                   shuffle=True)
 
-validation_dataset = LodopabDataset(cache_files=cache_files,
+validation_dataset = LodopabDataset(cache_files=small_cache_files,
                                     ground_truth_data_loc=ground_truth_data_loc,
-                                    split='validation')
+                                    split='validation',
+                                    transform=functools.partial(torch.unsqueeze, dim=0),
+                                    target_transform=functools.partial(torch.unsqueeze, dim=0))
 
 validation_dataloader = torch.utils.data.DataLoader(dataset=validation_dataset,
                                                     batch_size=16,
                                                     shuffle=True)
 
+# model defition
 model = UNet()
+print('model initialized:')
+summary(model)
+
+# training parameters
 criterion = torch.nn.MSELoss()
 optimizer = torch.optim.SGD(model.parameters(),
                             lr=0.01,
                             weight_decay=1e-8)
 
+# auto define the correct device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using {device} device.")
+
+# initiate Trainer
 trainer = Trainer(model=model,
                   device=torch.device(device),
                   criterion=criterion,
@@ -58,6 +66,11 @@ trainer = Trainer(model=model,
                   training_dataloader=training_dataloader,
                   validation_dataloader=validation_dataloader,
                   lr_scheduler=None,
-                  epochs=2,
+                  epochs=10,
                   epoch=0,
                   notebook=False)
+
+# start training
+training_losses, validation_losses, lr_rates = trainer.run_trainer()
+model_name = f'models/{date.today().strftime("%Y-%m-%d")}_unet.pt'
+torch.save(model.state_dict(), Path.cwd() / model_name)