text: import torch
  import torch.nn as nn
  import numpy as np
  import deeplab_resnet_2D
  import cv2
  from torch.autograd import Variable
  import torch.optim as optim
  import scipy.misc
  import torch.backends.cudnn as cudnn
  import sys
  import os
  from tqdm import *
  import random
  from docopt import docopt
  import sys
  
  sys.path.append('../')
  sys.path.append('../utils/')
  import PP
  #import matplotlib.pyplot as plt
  
  docstr = """Train ResNet-DeepLab on VOC12 (scenes) in pytorch using MSCOCO pretrained initialization 
  
  Usage: 
      train.py [options]
  
  Options:
      -h, --help                  Print this message
      --LISTpath=<str>            Input image number list file [default: ../../Data/MS2017a/train_slices.txt]
      --NoLabels=<int>            The number of different labels in training data, VOC has 21 labels, including background [default: 2]
      --lr=<float>                Learning Rate [default: 0.00025]
      -i, --iterSize=<int>        Num iters to accumulate gradients over [default: 10]
      --wtDecay=<float>           Weight decay during training [default: 0.0005]
      --gpu0=<int>                GPU number [default: 0]
      --maxIter=<int>             Maximum number of iterations [default: 20000]
      --useGPU=<int>              Use GPU or not [default: 0]
      --snapshotPath=<str>        Path to save snapshots [default: ../models/snapshots/]
      --snapshotPrefix=<str>      Snapshot model prefix [default: 2Ddeeplab_]
      --modelPath=<str>           Path of model to continue training on [default: none]
  """
  
  args = docopt(docstr, version='v0.1')
  print(args)
  
  gpu0 = int(args['--gpu0'])
  useGPU = int(args['--useGPU'])
  snapshot_path = args['--snapshotPath']
  snapshot_prefix = args['--snapshotPrefix']
  model_path = args['--modelPath']
  max_iter = int(args['--maxIter']) 
  batch_size = 1
  weight_decay = float(args['--wtDecay'])
  base_lr = float(args['--lr'])
  
  if int(args['--NoLabels']) == 2:
      onlyLesions = True
  else:
      onlyLesions = False
  
  if useGPU:
      cudnn.enabled = True
  else:
      cudnn.enabled = False
  
  
  def outS(i):
      """Given shape of input image as i,i,3 in deeplab-resnet model, this function
      returns j such that the shape of output blob of is j,j,21 """
      j = int(i)
      j = (j+1)/2
      j = int(np.ceil((j+1)/2.0))
      j = (j+1)/2
      return j
  
  def read_file(path_to_file):
      with open(path_to_file) as f:
          img_list = []
          for line in f:
              img_list.append(line[:-1])
      return img_list
  
  #returns a generator, returning the next chunk based on batch size
  def chunker(seq, size):
   return (seq[pos:pos+size] for pos in xrange(0,len(seq), size))
  
  def resize_label_batch(label, size):
      label_resized = np.zeros((size,size,1,label.shape[3]))
      interp = nn.UpsamplingBilinear2d(size=(size, size))
      labelVar = Variable(torch.from_numpy(label.transpose(3, 2, 0, 1)))
      label_resized[:, :, :, :] = interp(labelVar).data.numpy().transpose(2, 3, 1, 0)
  
      return label_resized
  
  def flip(I,flip_p):
      if flip_p>0.5:
          return np.fliplr(I)
      else:
          return I
  
  def scale_im(img_temp,scale):
      new_dims = (  int(img_temp.shape[0]*scale),  int(img_temp.shape[1]*scale)   )
      return cv2.resize(img_temp,new_dims).astype(float)
  
  def scale_gt(img_temp,scale):
      new_dims = (  int(img_temp.shape[0]*scale),  int(img_temp.shape[1]*scale)   )
      return cv2.resize(img_temp,new_dims,interpolation = cv2.INTER_NEAREST).astype(float)
     
  #Example chunk = ['scans/132/slices/FLAIR_28.nii.gz']
  def get_data_from_chunk_v2(chunk):
  
      main_folder_path = '../../Data/MS2017a/'
      scans_folder_path = main_folder_path + 'scans/'
  
      img_type_path = 'pre/FLAIR.nii.gz'
      gt_type_path = 'wmh.nii.gz'
  
      scale = random.uniform(0.5, 1.3)
      dim = int(scale*321)
  
      images = np.zeros((dim,dim, 1,len(chunk)))
      gt = np.zeros((dim,dim,1,len(chunk)))
      for i, piece in enumerate(chunk):
          print(os.path.join(main_folder_path, piece))
          img_temp = PP.numpyFromScan(os.path.join(main_folder_path, piece))
          flip_p = random.uniform(0, 1)
  
          img_temp = cv2.resize(img_temp,(321,321)).astype(float)
          img_temp = img_temp.reshape([321, 321, 1])
  
          img_temp = scale_im(img_temp,scale)
          img_temp = flip(img_temp,flip_p)
          images[:,:,0,i] = img_temp
  
          piece_gt = piece.replace('slices', 'gt_slices').replace('FLAIR', 'wmh')
          gt_temp = PP.numpyFromScan(os.path.join(main_folder_path, piece_gt), makebin = onlyLesions)
          gt_temp = cv2.resize(gt_temp,(321,321) , interpolation = cv2.INTER_NEAREST)
          gt_temp = gt_temp.reshape([321,321, 1])
          gt_temp = scale_gt(gt_temp,scale)
          gt_temp = flip(gt_temp,flip_p)
  
          gt[:,:,0,i] = gt_temp
          a = outS(321*scale)
  
      labels = [resize_label_batch(gt,i) for i in [a,a,a,a]]
  
      #from dim1 x dim2 x 1 x batch -> batch x 1 x dim1 x dim2
      images = images.transpose((3,2,0,1))
      images = torch.from_numpy(images).float()
      return images, labels
  
  def loss_calc(out, label, gpu0):
      """
      This function returns cross entropy loss for semantic segmentation
      """
      # out shape batch_size x channels x h x w -> batch_size x channels x h x w
      # label shape h x w x 1 x batch_size  -> batch_size x 1 x h x w
      label = label[:,:,0,:].transpose(2,0,1)
  
      label = torch.from_numpy(label).long()
      if useGPU:
          label = Variable(label).cuda(gpu0)
          if onlyLesions:
              criterion = nn.NLLLoss2d(weight = torch.cuda.FloatTensor([1, 100000]))
          else:
              criterion = nn.NLLLoss2d(weight = torch.cuda.FloatTensor([1, 100000, 100000]))
      else:
          label = Variable(label)
  
          if onlyLesions:
              criterion = nn.NLLLoss2d(weight = torch.FloatTensor([1, 100000]))
          else:
              criterion = nn.NLLLoss2d(weight = torch.FloatTensor([1, 100000, 100000]))
  
      m = nn.LogSoftmax()
      out = m(out)
  
      return criterion(out,label)
  
  
  def lr_poly(base_lr, iter,max_iter,power):
      return base_lr*((1-float(iter)/max_iter)**(power))
  
  def get_1x_lr_params_NOscale(model):
      """
      This generator returns all the parameters of the net except for 
      the last classification layer. Note that for each batchnorm layer, 
      requires_grad is set to False in deeplab_resnet.py, therefore this function does not return 
      any batchnorm parameter
      """
      b = []
  
      b.append(model.Scale.conv1)
      b.append(model.Scale.bn1)
      b.append(model.Scale.layer1)
      b.append(model.Scale.layer2)
      b.append(model.Scale.layer3)
      b.append(model.Scale.layer4)
  
      for i in range(len(b)):
          for j in b[i].modules():
              jj = 0
              for k in j.parameters():
                  jj+=1
                  if k.requires_grad:
                      yield k
  
  def get_10x_lr_params(model):
      """
      This generator returns all the parameters for the last layer of the net,
      which does the classification of pixel into classes
      """
      b = []
      b.append(model.Scale.layer5.parameters())
  
      for j in range(len(b)):
          for i in b[j]:
              yield i
  
  def modelInit():
      if not os.path.exists(snapshot_path):
          os.makedirs(snapshot_path)
  
      model = deeplab_resnet_2D.Res_Deeplab(int(args['--NoLabels']))
  
      if model_path != 'none':
          if useGPU:
              #loading on GPU when model was saved on GPU
              saved_state_dict = torch.load(model_path)
          else:
              #loading on CPU when model was saved on GPU
              saved_state_dict = torch.load(model_path, map_location=lambda storage, loc: storage)
          model.load_state_dict(saved_state_dict)
  
      model.float()
      model.eval() # use_global_stats = True
  
      return model
  
  
  def trainModel(model):
      img_list = read_file(args['--LISTpath'])
  
      data_list = []
      for i in range(50):  # make list for 10 epocs, though we will only use the first max_iter*batch_size entries of this list
          np.random.shuffle(img_list)
          data_list.extend(img_list)
  
      if useGPU:
          model.cuda(gpu0)
  
      #criterion = nn.CrossEntropyLoss() # use a Classification Cross-Entropy loss
      optimizer = optim.SGD([{'params': get_1x_lr_params_NOscale(model), 'lr': base_lr }, {'params': get_10x_lr_params(model), 'lr': 10*base_lr} ], lr = base_lr, momentum = 0.9,weight_decay = weight_decay)
  
      optimizer.zero_grad()
      data_gen = chunker(data_list, batch_size)
  
      for iter in range(max_iter+1):
          chunk = data_gen.next()
          images, label = get_data_from_chunk_v2(chunk)
  
          if useGPU:
              images = Variable(images).cuda(gpu0)
          else:
              images = Variable(images)
  
          out = model(images)
  
          loss = loss_calc(out[0], label[0],gpu0)
          iter_size = int(args['--iterSize']) 
          for i in range(len(out)-1):
              loss = loss + loss_calc(out[i+1],label[i+1],gpu0)
          loss = loss/iter_size 
          loss.backward()
  
          if iter % 1 == 0:
              print 'iter = ',iter, 'of',max_iter,'completed, loss = ', iter_size*(loss.data.cpu().numpy())
  
          if iter % iter_size  == 0:
              optimizer.step()
              lr_ = lr_poly(base_lr,iter,max_iter,0.9)
              print '(poly lr policy) learning rate',lr_
              optimizer = optim.SGD([{'params': get_1x_lr_params_NOscale(model), 'lr': lr_ }, {'params': get_10x_lr_params(model), 'lr': 10*lr_} ], lr = lr_, momentum = 0.9,weight_decay = weight_decay)
              optimizer.zero_grad()
  
          if iter % 1000 == 0 and iter!=0:
              print 'taking snapshot ...'
              torch.save(model.state_dict(), os.path.join(snapshot_path, snapshot_prefix + str(iter) + '.pth'))
  
  if __name__ == "__main__":
      model = modelInit()
      trainModel(model)    