text: import argparse
  
  import numpy as np
  from torch.utils.data import DataLoader
  from torchvision import transforms
  
  from dataset import CustomDataset
  
  parser = argparse.ArgumentParser(description='Train the UNet on images and target masks',
                                       formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  parser.add_argument('-e', '--epochs', metavar='E', type=int, default=5,
                          help='Number of epochs', dest='epochs')
  parser.add_argument('-b', '--batch-size', metavar='B', type=int, nargs='?', default=1,
                          help='Batch size', dest='batchsize')
  parser.add_argument('-l', '--learning-rate', metavar='LR', type=float, nargs='?', default=0.0001,
                          help='Learning rate', dest='lr')
  parser.add_argument('-f', '--load', dest='load', type=str, default=False,
                          help='Load model from a .pth file')
  parser.add_argument('-s', '--scale', dest='scale', type=float, default=0.5,
                          help='Downscaling factor of the images')
  parser.add_argument('-v', '--validation', dest='val', type=float, default=10.0,
                          help='Percent of the data that is used as validation (0-100)')
  
  parser.add_argument('--n_classes', type=int, default=1)
  
  args = parser.parse_args()
  
  
  # use same transform for train/val for this example
  trans = transforms.Compose([
      transforms.ToTensor(),
      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # imagenet
  ])
  
  # train_set = SimDataset(200, transform=trans)
  # val_set = SimDataset(200, transform=trans)
  
  num_class = args.n_classes
  n_epochs = args.epochs
  batch_size = args.batchsize
  
  train_set = CustomDataset(num_class, 'data/train/imgs', 'data/train/masks', transform=trans)
  val_set = CustomDataset(num_class, 'data/val/imgs', 'data/val/masks', transform=trans)
  
  image_datasets = {
      'train': train_set, 'val': val_set
  }
  
  dataloaders = {
      'train': DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0),
      'val': DataLoader(val_set, batch_size=batch_size, shuffle=True, num_workers=0)
  }
  
  dataset_sizes = {
      x: len(image_datasets[x]) for x in image_datasets.keys()
  }
  
  print(dataset_sizes)
  
  def reverse_transform(inp):
      inp = inp.numpy().transpose((1, 2, 0))
      mean = np.array([0.485, 0.456, 0.406])
      std = np.array([0.229, 0.224, 0.225])
      inp = std * inp + mean
      inp = np.clip(inp, 0, 1)
      inp = (inp * 255).astype(np.uint8)
  
      return inp
  
  
  # Get a batch of training data
  inputs, masks = next(iter(dataloaders['train']))
  
  print(inputs.shape, masks.shape)
  for x in [inputs.numpy(), masks.numpy()]:
      print(x.min(), x.max(), x.mean(), x.std())
  
  # plt.imshow(reverse_transform(inputs[3]))
  
  from collections import defaultdict
  import torch.nn.functional as F
  from loss import dice_loss
  
  
  def calc_loss(pred, target, metrics, bce_weight=0.5):
      bce = F.binary_cross_entropy_with_logits(pred, target)
  
      pred = torch.sigmoid(pred)
      dice = dice_loss(pred, target)
  
      loss = bce * bce_weight + dice * (1 - bce_weight)
  
      metrics['bce'] += bce.data.cpu().numpy() * target.size(0)
      metrics['dice'] += dice.data.cpu().numpy() * target.size(0)
      metrics['loss'] += loss.data.cpu().numpy() * target.size(0)
  
      return loss
  
  
  def print_metrics(metrics, epoch_samples, phase):
      outputs = []
      for k in metrics.keys():
          outputs.append("{}: {:4f}".format(k, metrics[k] / epoch_samples))
  
      print("{}: {}".format(phase, ", ".join(outputs)))
  
  
  def train_model(model, optimizer, scheduler, num_epochs=25):
      best_model_wts = copy.deepcopy(model.state_dict())
      best_loss = 1e10
  
      for epoch in range(num_epochs):
          print('Epoch {}/{}'.format(epoch, num_epochs - 1))
          print('-' * 10)
  
          since = time.time()
  
          # Each epoch has a training and validation phase
          for phase in ['train', 'val']:
              if phase == 'train':
                  scheduler.step()
                  for param_group in optimizer.param_groups:
                      print("LR", param_group['lr'])
  
                  model.train()  # Set model to training mode
              else:
                  model.eval()  # Set model to evaluate mode
  
              metrics = defaultdict(float)
              epoch_samples = 0
  
              for inputs, labels in dataloaders[phase]:
                  inputs = inputs.to(device)
                  labels = labels.to(device)
  
                  # zero the parameter gradients
                  optimizer.zero_grad()
  
                  # forward
                  # track history if only in train
                  with torch.set_grad_enabled(phase == 'train'):
                      outputs = model(inputs)
                      loss = calc_loss(outputs, labels, metrics)
  
                      # backward + optimize only if in training phase
                      if phase == 'train':
                          loss.backward()
                          optimizer.step()
  
                  # statistics
                  epoch_samples += inputs.size(0)
  
              print_metrics(metrics, epoch_samples, phase)
              epoch_loss = metrics['loss'] / epoch_samples
  
              # deep copy the model
              if phase == 'val' and epoch_loss < best_loss:
                  print("saving best model")
                  best_loss = epoch_loss
                  best_model_wts = copy.deepcopy(model.state_dict())
  
          time_elapsed = time.time() - since
          print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
  
      print('Best val loss: {:4f}'.format(best_loss))
  
      # load best model weights
      model.load_state_dict(best_model_wts)
      return model
  
  
  import torch
  import torch.optim as optim
  from torch.optim import lr_scheduler
  import time
  import copy
  from resnet_unet import ResNetUNet
  
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  print(device)
  
  model = ResNetUNet(num_class).to(device)
  
  # freeze backbone layers
  # Comment out to finetune further
  for l in model.base_layers:
      for param in l.parameters():
          param.requires_grad = False
  
  optimizer_ft = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-4)
  
  exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
  
  model = train_model(model, optimizer_ft, exp_lr_scheduler, num_epochs=n_epochs)
  
  torch.save(model.state_dict(), 'best.pth')