1. added generic dataset loader

2. fix driver.py timesteps=args.timesteps instead of args.epochs
3. add save model in driver.py
4. add load model in driver.py
5. fix wandb image logger that requires channels last or no channels.
6. added ability to evaluate and save model every n epochs to train faster.

Author: Aaron

driver.py

@@ -1,56 +1,69 @@
 import os
 import argparse
 from tqdm import tqdm
-
 import torch
+import numpy as np
 import torchvision.transforms as transforms
-
 from torch.optim import AdamW
 from lion_pytorch import Lion
-
 from med_seg_diff_pytorch import Unet, MedSegDiff
-from med_seg_diff_pytorch.dataset import ISICDataset
-
+from med_seg_diff_pytorch.dataset import ISICDataset, GenericNpyDataset
 from accelerate import Accelerator
 import wandb
 
 ## Parse CLI arguments ##
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('-slr', '--scale_lr', action='store_true', help="Whether to scale lr.")
-    parser.add_argument('-rt', '--report_to', type=str, default="wandb", choices=["wandb"], help="Where to log to. Currently only supports wandb")
+    parser.add_argument('-rt', '--report_to', type=str, default="wandb", choices=["wandb"],
+                        help="Where to log to. Currently only supports wandb")
     parser.add_argument('-ld', '--logging_dir', type=str, default="logs", help="Logging dir.")
     parser.add_argument('-od', '--output_dir', type=str, default="output", help="Output dir.")
-    parser.add_argument('-mp', '--mixed_precision', type=str, default="no", choices=["no", "fp16", "bf16"], help="Whether to do mixed precision")
-    parser.add_argument('-ga', '--gradient_accumulation_steps', type=int, default=4, help="The number of gradient accumulation steps.")
-    parser.add_argument('-img', '--img_folder', type=str, default='ISBI2016_ISIC_Part3B_Training_Data', help='The image file path from data_path')
-    parser.add_argument('-csv', '--csv_file', type=str, default='ISBI2016_ISIC_Part3B_Training_GroundTruth.csv', help='The csv file to load in from data_path')
+    parser.add_argument('-mp', '--mixed_precision', type=str, default="no", choices=["no", "fp16", "bf16"],
+                        help="Whether to do mixed precision")
+    parser.add_argument('-ga', '--gradient_accumulation_steps', type=int, default=4,
+                        help="The number of gradient accumulation steps.")
+    parser.add_argument('-img', '--img_folder', type=str, default='ISBI2016_ISIC_Part3B_Training_Data',
+                        help='The image file path from data_path')
+    parser.add_argument('-csv', '--csv_file', type=str, default='ISBI2016_ISIC_Part3B_Training_GroundTruth.csv',
+                        help='The csv file to load in from data_path')
     parser.add_argument('-sc', '--self_condition', action='store_true', help='Whether to do self condition')
     parser.add_argument('-lr', '--learning_rate', type=float, default=5e-4, help='learning rate')
-    parser.add_argument('-ab1', '--adam_beta1', type=float, default=0.95, help='The beta1 parameter for the Adam optimizer.')
-    parser.add_argument('-ab2', '--adam_beta2', type=float, default=0.999, help='The beta2 parameter for the Adam optimizer.')
-    parser.add_argument('-aw', '--adam_weight_decay', type=float, default=1e-6, help='Weight decay magnitude for the Adam optimizer.')
-    parser.add_argument('-ae', '--adam_epsilon', type=float, default=1e-08, help='Epsilon value for the Adam optimizer.')
+    parser.add_argument('-ab1', '--adam_beta1', type=float, default=0.95,
+                        help='The beta1 parameter for the Adam optimizer.')
+    parser.add_argument('-ab2', '--adam_beta2', type=float, default=0.999,
+                        help='The beta2 parameter for the Adam optimizer.')
+    parser.add_argument('-aw', '--adam_weight_decay', type=float, default=1e-6,
+                        help='Weight decay magnitude for the Adam optimizer.')
+    parser.add_argument('-ae', '--adam_epsilon', type=float, default=1e-08,
+                        help='Epsilon value for the Adam optimizer.')
     parser.add_argument('-ul', '--use_lion', type=bool, default=False, help='use Lion optimizer')
     parser.add_argument('-ic', '--mask_channels', type=int, default=1, help='input channels for training (default: 3)')
-    parser.add_argument('-c', '--input_img_channels', type=int, default=3, help='output channels for training (default: 3)')
+    parser.add_argument('-c', '--input_img_channels', type=int, default=3,
+                        help='output channels for training (default: 3)')
     parser.add_argument('-is', '--image_size', type=int, default=128, help='input image size (default: 128)')
     parser.add_argument('-dd', '--data_path', default='./data', help='directory of input image')
-    parser.add_argument('-d', '--dim', type=int, default=64, help='dim (deaault: 64)')
-    parser.add_argument('-e', '--epochs', type=int, default=10, help='number of epochs (default: 128)')
+    parser.add_argument('-d', '--dim', type=int, default=64, help='dim (default: 64)')
+    parser.add_argument('-e', '--epochs', type=int, default=10000, help='number of epochs (default: 10000)')
     parser.add_argument('-bs', '--batch_size', type=int, default=8, help='batch size to train on (default: 8)')
-    parser.add_argument('-ds', '--dataset', default='ISIC', help='Dataset to use')
+    parser.add_argument('--timesteps', type=int, default=1000, help='number of timesteps (default: 1000)')
+    parser.add_argument('-ds', '--dataset', default='generic', help='Dataset to use')
+    parser.add_argument('--save_every', type=int, default=100, help='save_every n rpochs (default: 100)')
+    parser.add_argument('--load_model_from', default=None, help='path to pt file to load from')
     return parser.parse_args()
 
 
 def load_data(args):
-    # Create transforms for data
-    transform_list = [transforms.Resize((args.image_size,args.image_size)), transforms.ToTensor(),]
-    transform_train = transforms.Compose(transform_list)
-
     # Load dataset
     if args.dataset == 'ISIC':
-        dataset = ISICDataset(args.data_path, args.csv_file, args.img_folder, transform = transform_train, training = True, flip_p=0.5)
+        transform_list = [transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), ]
+        transform_train = transforms.Compose(transform_list)
+        dataset = ISICDataset(args.data_path, args.csv_file, args.img_folder, transform=transform_train, training=True,
+                              flip_p=0.5)
+    elif args.dataset == 'generic':
+        transform_list = [transforms.ToPILImage(), transforms.Resize(args.image_size), transforms.ToTensor()]
+        transform_train = transforms.Compose(transform_list)
+        dataset = GenericNpyDataset(args.data_path, transform=transform_train, test_flag=False)
     else:
         raise NotImplementedError(f"Your dataset {args.dataset} hasn't been implemented yet.")
 
@@ -63,10 +76,11 @@ def load_data(args):
     return training_generator
 
 
-
 def main():
     args = parse_args()
+    checkpoint_dir = os.path.join(args.output_dir, 'checkpoints')
     logging_dir = os.path.join(args.output_dir, args.logging_dir)
+    os.makedirs(checkpoint_dir, exist_ok=True)
     accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
@@ -78,20 +92,20 @@ def main():
 
     ## DEFINE MODEL ##
     model = Unet(
-        dim = args.dim,
-        image_size = args.image_size,
-        dim_mults = (1, 2, 4, 8),
-        mask_channels = args.mask_channels,
-        input_img_channels= args.input_img_channels,
-        self_condition = args.self_condition
+        dim=args.dim,
+        image_size=args.image_size,
+        dim_mults=(1, 2, 4, 8),
+        mask_channels=args.mask_channels,
+        input_img_channels=args.input_img_channels,
+        self_condition=args.self_condition
     )
 
     ## LOAD DATA ##
     data_loader = load_data(args)
-    #training_generator = tqdm(data_loader, total=int(len(data_loader)))
+    # training_generator = tqdm(data_loader, total=int(len(data_loader)))
     if args.scale_lr:
         args.learning_rate = (
-            args.learning_rate * args.gradient_accumulation_steps * args.batch_size * accelerator.num_processes
+                args.learning_rate * args.gradient_accumulation_steps * args.batch_size * accelerator.num_processes
         )
 
     ## Initialize optimizer
@@ -112,22 +126,29 @@ def main():
         )
 
     ## TRAIN MODEL ##
-    running_loss = 0.0
     counter = 0
     model, optimizer, data_loader = accelerator.prepare(
         model, optimizer, data_loader
     )
     diffusion = MedSegDiff(
         model,
-        timesteps = args.epochs
+        timesteps=args.timesteps
     ).to(accelerator.device)
+
+    if args.load_model_from is not None:
+        save_dict = torch.load(args.load_model_from)
+        diffusion.model.load_state_dict(save_dict['model_state_dict'])
+        optimizer.load_state_dict(save_dict['optimizer_state_dict'])
+        accelerator.print(f'Loaded from {args.load_model_from}')
+
     ## Iterate across training loop
     for epoch in range(args.epochs):
-        print('Epoch {}/{}'.format(epoch+1, args.epochs))
+        running_loss = 0.0
+        print('Epoch {}/{}'.format(epoch + 1, args.epochs))
         for (img, mask) in tqdm(data_loader):
             with accelerator.accumulate(model):
                 loss = diffusion(mask, img)
-                accelerator.log({'loss': loss}) # Log loss to wandb
+                accelerator.log({'loss': loss})  # Log loss to wandb
                 accelerator.backward(loss)
                 optimizer.step()
                 optimizer.zero_grad()
@@ -136,12 +157,24 @@ def main():
         epoch_loss = running_loss / len(data_loader)
         print('Training Loss : {:.4f}'.format(epoch_loss))
         ## INFERENCE ##
-        pred = diffusion.sample(img).cpu().detach().numpy()
-        for tracker in accelerator.trackers:
-            if tracker.name == "wandb":
-                tracker.log(
-                    {'pred-img-mask': [wandb.Image(pred), wandb.Image(img), wandb.Image(mask)]}
-                )
+
+        if epoch % args.save_every == 0:
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': diffusion.model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'loss': loss,
+            }, os.path.join(checkpoint_dir, f'state_dict_epoch_{epoch}_loss_{epoch_loss}.pt'))
+
+            pred = diffusion.sample(img).cpu().detach().numpy()
+
+            for tracker in accelerator.trackers:
+                if tracker.name == "wandb":
+                    # save just one image per batch
+                    tracker.log(
+                        {'pred-img-mask': [wandb.Image(pred[0, 0, :, :]), wandb.Image(img[0, 0, :, :]),
+                                           wandb.Image(mask[0, 0, :, :])]}
+                    )
 
 
 if __name__ == '__main__':

med_seg_diff_pytorch/dataset.py

@@ -1,26 +1,31 @@
 import os
-os.environ['KMP_DUPLICATE_LIB_OK']='True'
+import numpy as np
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 import torch
 from torch.utils.data import Dataset
 from PIL import Image
 import pandas as pd
 import random
 import torchvision.transforms.functional as F
+
+
 class ISICDataset(Dataset):
-    def __init__(self, data_path, csv_file, img_folder, transform = None, training = True, flip_p=0.5):
+    def __init__(self, data_path, csv_file, img_folder, transform=None, training=True, flip_p=0.5):
         df = pd.read_csv(os.path.join(data_path, csv_file), encoding='gbk')
         self.img_folder = img_folder
-        self.name_list = df.iloc[:,0].tolist()
-        self.label_list = df.iloc[:,1].tolist()
+        self.name_list = df.iloc[:, 0].tolist()
+        self.label_list = df.iloc[:, 1].tolist()
         self.data_path = data_path
         self.transform = transform
         self.training = training
         self.flip_p = flip_p
+
     def __len__(self):
         return len(self.name_list)
+
     def __getitem__(self, index):
         """Get the images"""
-        name = self.name_list[index]+'.jpg'
+        name = self.name_list[index] + '.jpg'
         img_path = os.path.join(self.data_path, self.img_folder, name)
 
         mask_name = name.split('.')[0] + '_Segmentation.png'
@@ -35,13 +40,52 @@ def __getitem__(self, index):
             label = int(self.label_list[index])
 
         if self.transform:
+            # save random state so that if more elaborate transforms are used
+            # the same transform will be applied to both the mask and the img
+            state = torch.get_rng_state()
             img = self.transform(img)
+            torch.set_rng_state(state)
             mask = self.transform(mask)
             if random.random() < self.flip_p:
                 img = F.vflip(img)
                 mask = F.vflip(mask)
 
-
         if self.training:
             return (img, mask)
         return (img, mask, label)
+
+
+class GenericNpyDataset(torch.utils.data.Dataset):
+    def __init__(self, directory: str, transform, test_flag: bool = True):
+        '''
+        Genereic dataset for loading npy files.
+        The npy store 3D arrays with the first two dimensions being the image and the third dimension being the channels.
+        channel 0 is the image and the other channel is the label.
+        '''
+        super().__init__()
+        self.directory = os.path.expanduser(directory)
+        self.transform = transform
+        self.test_flag = test_flag
+        self.filenames = [os.path.join(self.directory, x) for x in os.listdir(self.directory) if x.endswith('.npy')]
+
+    def __getitem__(self, x: int):
+        fname = self.filenames[x]
+        npy_img = np.load(fname)
+        img = npy_img[:, :, :1]
+        img = torch.from_numpy(img).permute(2, 0, 1)
+        mask = npy_img[:, :, 1:]
+        mask = np.where(mask > 0, 1, 0)
+        image = img[:, ...]
+        mask = torch.from_numpy(mask).permute(2, 0, 1).float()
+        if self.transform:
+            # save random state so that if more elaborate transforms are used
+            # the same transform will be applied to both the mask and the img
+            state = torch.get_rng_state()
+            image = self.transform(image)
+            torch.set_rng_state(state)
+            mask = self.transform(mask)
+
+        return image, mask
+
+    def __len__(self) -> int:
+        return len(self.filenames)