Merge pull request #7 from bf2harven/main

fixes

Author: lucidrains

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)