<jupyter_start><jupyter_text>"Training a 3D Unet to perform brain tumor segmentation (BraTS)"> "in which we train a Unet using Colab's free GPUs"- toc: false- branch: master- badges: true- comments: true- categories: [fastpages, jupyter, segmentation]- image: images/some_folder/your_image.png- hide: false- search_exclude: false- metadata_key1: metadata_value1- metadata_key2: metadata_value2   A compact introduction to 3D Unets for medical image segmentation.For the past few months (since late January) I have been running a deep learning course for doctors (mainly aimed at (neuro)radiologists, neurosurgeons, psychiatrists, and other MDs who work with imaging on a daily basis).  This course is part of the Certificate of Advanced Studies in Medical Imaging AI that we launched in conjunction with the Swiss Institute for Translational and Entrepreneurial Medicine last September: > twitter: https://twitter.com/mckinley_scan/status/1124297078857969664 We've primarily been working with the [fastai](www.fast.ai) library, for three main reasons.  Number one, I'm a big fan of Pytorch as a deep learning library, and fastai is built on top of Pytorch.  Number two, fastai comes with an awesome online course which several of our students have already tapped into as an additional resource.  Finally, fastai makes things like learning rate tuning, and transfer learning really easy.  One of the primary problems we looked at in our course is medical image segmentation (it's also the topic I'm most heavily involved with in my research).  Fastai has a really nice feature for image segmentation called [Dynamic Unet](https://medium.com/@keremturgutlu/semantic-segmentation-u-net-part-1-d8d6f6005066), which makes it easy to leverage pretrained image *classifiers* for image segmentation.  We've (meaning both myself and the students) been really impressed by how well Dynamic Unets perform on medical imaging tasks.  One thing fastai doesn't do (yet?) is provide easy access to 3D segmentation models: to segment a 3D volume using fastai, we have to slice it into 2D images and then segment slice-by-slice.  It's natural to ask if models which can incorporate 3D context perform substantially better than purely 2D models.  While there are a huge number of models and model variations out there for volumetric segmentation, the research group at the DKFZ have had substantial success in the last few years applying a rather normal 3D Unet variant to a variety of segmentation tasks, an approach they call [no-new-Unet](https://arxiv.org/abs/1809.10486).    This post/notebookThroughout the course, I have been providing Jupyter notebooks to our students illustrating image classification and segmentation, using fastai and Google Colab.  This approach has a number of advantages: Jupyter allows text, images and code to be freely mixed, and Colab allows students to run code without installing python on their own machine.  Colab even allows access to pretty beefy GPUs, if only for a limited amount of time, and this led me to the current post.  A version of nnUnet took 2nd place in the 2018 BraTS challenge for Brain Tumor segmentation, losing out only to an entry from NVIDIA which used 32Gb GPUs to scale out their neural network far beyond what most participants could hope for.  (Our own entry, DeepSCAN, was joint third that year, of which more in a future post).  You can download nnUnet and train your own BraTS model: on my machine, with a fairly modern GPU, it takes 4-5 days for the network to conclude training, which is far in excess of what you can achieve in Colab.  This post is an attempt to fit the training of a 3D Unet into Colab, adhering roughly to how training works in fastai (learning rate finding, one-cycle learning rate schedule).It *sort-of* works.  Each 'epoch' of 1000 batches takes about half an hour, on a high RAM Colab instance (which has 4 CPU cores).  Without babysitting your colab session or employing *nefarious* tactics to keep your session running, you'll only manage to train and validate a couple of epochs.  Still, I think it's interesting and worth a try (especially as the main developer (Hi Fabian!) of nnUnet is apparently morally opposed to Jupyter Notebooks)If you want to try it yourself, this blog post was generated from a Jupyter notebook (using [fastpages](https://fastpages.fast.ai/) that you can load into colab using the badge above.    The DataThe BraTS dataset is a collection of multi-modal Magnetic Resonance images of patients with a low- or high-grade glioma, derived from several hospitals, and the<jupyter_code>%%capture 
#hide

!mkdir BRATS_dkfz
!wget http://mia-software.artorg.unibe.ch/CAS_zips/dkfz_uint.zip 
!unzip 'dkfz_uint.zip' -d BRATS_dkfz

!rm dkfz_uint.zip<jupyter_output><empty_output><jupyter_text>We need some packages: batchgenerators (to build an efficient dataloader for 3D data with augmentations) and nnunet (to get a modern implementation of a 3D Unet, with patchwise application)<jupyter_code>%%capture 
#hide

!pip install --upgrade batchgenerators
!pip install SimpleITK
!pip install nnunet
#hide

import multiprocessing

multiprocessing.cpu_count()<jupyter_output><empty_output><jupyter_text>Creating a dataloader<jupyter_code>#hide

#set this to True, if you want all the 3D image files to be loaded into memory
IN_MEMORY = True

#set this above zero to limit the number of training cases 
USE_TRAIN_CASES = 0

#set this above zero to limit the number of validation cases (to make epochs faster)
USE_VAL_CASES = 0

#data loader/augmentation settings
MAX_ROT_ANGLE = 15

P_ROT = 0.1

P_SCALE = 0.1

P_EL = 0.1

MAX_SHAPE = (192,192,192)

NUM_THREADS = 4

#training settings

NUM_EPOCHS = 2

BATCHES_PER_EPOCH = 1000

DICE_WEIGHTING = 1.0

OUTPUT_DIR = "brats_2018_local"

PATCH_SIZE = (128, 128, 128)

BATCH_SIZE = 2

GRADIENT_MASK_ZERO_VOXELS = False



TARGET_LABEL_SETS = [ [3],
                     [1,3],
                     [1,2,3]
                    ]
TARGET_LABEL_NAMES = ['enhancing', 'tumor_core', 'whole_tumor']

#hide


%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt

import numpy as np


from time import time
from batchgenerators.augmentations.crop_and_pad_augmentations import crop
from batchgenerators.dataloading import MultiThreadedAugmenter
from batchgenerators.transforms import Compose
from batchgenerators.utilities.data_splitting import get_split_deterministic
from batchgenerators.utilities.file_and_folder_operations import *
from batchgenerators.dataloading.data_loader import DataLoader
from batchgenerators.augmentations.utils import pad_nd_image
from batchgenerators.transforms.spatial_transforms import SpatialTransform_2, MirrorTransform
from batchgenerators.transforms.color_transforms import BrightnessMultiplicativeTransform, GammaTransform
from batchgenerators.transforms.noise_transforms import GaussianNoiseTransform, GaussianBlurTransform


import os
import random
from tqdm.notebook import tqdm

#torch-related imports

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from torch.autograd import Variable, Function
#from torch.nn import BCELoss, NLLLoss, BCEWithLogitsLoss, MSELoss, ModuleList, ReplicationPad2d
#from torch.utils.data.sampler import Sampler, RandomSampler, WeightedRandomSampler, SequentialSampler
#from torch._six import int_classes as _int_classes


from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR

from nnunet.network_architecture.generic_UNet import Generic_UNet




#hide

brats_preprocessed_folder = '/content/BRATS_dkfz'

def get_list_of_patients(preprocessed_data_folder):
    npy_files = subfiles(preprocessed_data_folder, suffix=".npy", join=True)
    # remove npy file extension
    patients = [i[:-4] for i in npy_files]
    return patients

patients = get_list_of_patients(brats_preprocessed_folder)

train, val = get_split_deterministic(patients, fold=0, num_splits=5, random_state=12345)

if USE_TRAIN_CASES != 0:
  train = train[:USE_TRAIN_CASES]

if USE_VAL_CASES != 0:
  val = val[:USE_VAL_CASES]

num_val_cases = len(val)

num_train_cases = len(train)

if IN_MEMORY:
    patient_volumes = {x:np.load(x+'.npy') for x in patients}
else:
    patient_volumes = {x:np.load(x+'.npy', 'r') for x in patients}

 #hide
 
 def load_patient(patient, axis = 'random'):
        if axis=='random':
          axis = np.random.choice(['axial', 'sagittal', 'coronal'])
        
        assert axis in ['axial', 'sagittal', 'coronal']


        data = patient_volumes[patient]

        if axis == 'coronal':
          data = np.swapaxes(data, 1,2)

        if axis == 'sagittal':
          data = np.swapaxes(data, 1,3)


        metadata = load_pickle(patient + ".pkl")

        return data, metadata

class BraTS2017DataLoader3D(DataLoader):
    def __init__(self, data, batch_size, patch_size, num_threads_in_multithreaded, seed_for_shuffle=1234,
                 return_incomplete=False, shuffle=True, infinite=True):
        """
        data must be a list of patients as returned by get_list_of_patients (and split by get_split_deterministic)
        patch_size is the spatial size the retured batch will have
        """
        super().__init__(data, batch_size, num_threads_in_multithreaded, seed_for_shuffle, return_incomplete, shuffle,
                         infinite)
        self.patch_size = patch_size
        self.num_modalities = 4
        self.indices = list(range(len(data)))

    @staticmethod
    def load_patient(patient, axis = 'random'):
        if axis=='random':
          axis = np.random.choice(['axial', 'sagittal', 'coronal'])
        
        assert axis in ['axial', 'sagittal', 'coronal']


        data = patient_volumes[patient]

        if axis == 'coronal':
          data = np.swapaxes(data, 1,2)

        if axis == 'sagittal':
          data = np.swapaxes(data, 1,3)
        metadata = load_pickle(patient + ".pkl")

        return data, metadata
   

    def generate_train_batch(self):
        # DataLoader has its own methods for selecting what patients to use next, see its Documentation
        idx = self.get_indices()
        patients_for_batch = [self._data[i] for i in idx]

        # initialize empty array for data and seg
        data = np.zeros((self.batch_size, self.num_modalities, *self.patch_size), dtype=np.float32)
        seg = np.zeros((self.batch_size, 1, *self.patch_size), dtype=np.float32)

        metadata = []
        patient_names = []

        # iterate over patients_for_batch and include them in the batch
        for i, j in enumerate(patients_for_batch):
            patient_data, patient_metadata = load_patient(j)

            patient_data = patient_data.astype(np.float)

            # this will only pad patient_data if its shape is smaller than self.patch_size
            patient_data = pad_nd_image(patient_data, self.patch_size)

            # now random crop to self.patch_size
            # crop expects the data to be (b, c, x, y, z) but patient_data is (c, x, y, z) so we need to add one
            # dummy dimension in order for it to work (@Todo, could be improved)
            patient_data, patient_seg = crop(patient_data[:-1][None], patient_data[-1:][None], self.patch_size, crop_type="random")

            zero_mask = patient_data == 0
            patient_data = patient_data - np.array(patient_metadata['means'])[:,np.newaxis,np.newaxis,np.newaxis]
            patient_data = patient_data/np.array(patient_metadata['sds'])[:,np.newaxis,np.newaxis,np.newaxis]

            patient_data[zero_mask] = 0.0



            # find all voxels in which each modality is zero
            zero_mask_all = np.all(patient_data[0]==0.0, axis=0)

            #code the zeros in the segmentation mask as a negative number

            patient_seg[zero_mask_all[np.newaxis,np.newaxis,:]] = -1

            data[i] = patient_data[0]
            seg[i] = patient_seg[0]

            metadata.append(patient_metadata)
            patient_names.append(j)

        return {'data': data, 'seg':seg, 'metadata':metadata, 'names':patient_names, }

#hide

def get_train_transform(patch_size):
    # we now create a list of transforms. These are not necessarily the best transforms to use for BraTS, this is just
    # to showcase some things
    tr_transforms = []

    # the first thing we want to run is the SpatialTransform. It reduces the size of our data to patch_size and thus
    # also reduces the computational cost of all subsequent operations. All subsequent operations do not modify the
    # shape and do not transform spatially, so no border artifacts will be introduced
    # Here we use the new SpatialTransform_2 which uses a new way of parameterizing elastic_deform
    # We use all spatial transformations with a probability of 0.2 per sample. This means that 1 - (1 - 0.1) ** 3 = 27%
    # of samples will be augmented, the rest will just be cropped
    tr_transforms.append(
        SpatialTransform_2(
            patch_size, [i // 2 for i in patch_size],
            do_elastic_deform=True, deformation_scale=(0, 0.25),
            do_rotation=True,
            angle_x=(- MAX_ROT_ANGLE  / 360. * 2 * np.pi, MAX_ROT_ANGLE  / 360. * 2 * np.pi),
            angle_y=(- MAX_ROT_ANGLE  / 360. * 2 * np.pi, MAX_ROT_ANGLE  / 360. * 2 * np.pi),
            angle_z=(- MAX_ROT_ANGLE  / 360. * 2 * np.pi, MAX_ROT_ANGLE  / 360. * 2 * np.pi),
            do_scale=True, scale=(0.75, 1.25),
            border_mode_data='constant', border_cval_data=0,
            border_mode_seg='constant', border_cval_seg=0,
            order_seg=1, order_data=3,
            random_crop=True,
            p_el_per_sample=P_EL, p_rot_per_sample=P_ROT, p_scale_per_sample=P_SCALE
        )
    )

    # now we mirror along all axes
    tr_transforms.append(MirrorTransform(axes=(0, 1, 2)))

    # brightness transform for 15% of samples
    tr_transforms.append(BrightnessMultiplicativeTransform((0.7, 1.5), per_channel=True, p_per_sample=0.5))

    # gamma transform. This is a nonlinear transformation of intensity values
    # (https://en.wikipedia.org/wiki/Gamma_correction)
    tr_transforms.append(GammaTransform(gamma_range=(0.5, 2), invert_image=False, per_channel=True, p_per_sample=0.25))
    
  
    # Gaussian Noise
    tr_transforms.append(GaussianNoiseTransform(noise_variance=(0, 0.05), p_per_sample=0.25))

    # blurring. Some BraTS cases have very blurry modalities. This can simulate more patients with this problem and
    # thus make the model more robust to it
    tr_transforms.append(GaussianBlurTransform(blur_sigma=(0.5, 1.5), different_sigma_per_channel=True,
                                               p_per_channel=0.5, p_per_sample=0.15))

    # now we compose these transforms together
    tr_transforms = Compose(tr_transforms)
    return tr_transforms
#hide 

dataloader_train = BraTS2017DataLoader3D(train, BATCH_SIZE, MAX_SHAPE, NUM_THREADS)

tr_transforms = get_train_transform(PATCH_SIZE)

# finally we can create multithreaded transforms that we can actually use for training
# we don't pin memory here because this is pytorch specific.
tr_gen = MultiThreadedAugmenter(dataloader_train, tr_transforms, num_processes=NUM_THREADS,
                                num_cached_per_queue=20,
                                seeds=None, pin_memory=True)
#


tr_gen.restart()

#hide_input


plt.rcParams["figure.figsize"] = (15,7)


batch = tr_gen.next()
image_data = batch['data']
seg_data = batch['seg']



zero_mask = seg_data[:,0]==-1

image_data = image_data*(1-zero_mask)[:,np.newaxis]

plt.imshow(np.concatenate([image_data[1,0, 64],
                           image_data[1,1, 64],
                           image_data[1,2, 64],
                           image_data[1,3, 64]], 1), cmap='gray')


def get_gradient_mask(batch):
  if GRADIENT_MASK_ZERO_VOXELS:
    return (batch[:,0] != 0.0).astype(np.float)
  else:
    return (np.ones_like(batch[:,0])).astype(np.float)


assert len(TARGET_LABEL_NAMES) == len(TARGET_LABEL_SETS)




total_with_gradient = 0

total_positives = np.array([0,0,0])

for case in train:


  total_with_gradient += np.sum(get_gradient_mask(patient_volumes[case][np.newaxis]))

  total_positives += [np.sum(np.isin(patient_volumes[case][-1], x)) for x in TARGET_LABEL_SETS]

label_prevalences = total_positives/total_with_gradient
list(zip(label_prevalences,TARGET_LABEL_NAMES))

def make_target(gt):
  return np.concatenate([np.isin(gt,labelset).astype(np.float) for labelset in TARGET_LABEL_SETS], axis =1)

class BCE_from_logits(nn.modules.Module):
    def __init__(self):
        super(BCE_from_logits,self).__init__()
    def forward(self, input, target):
        #input = input.clamp(min = -1, max = 1)
        max_val = (-input).clamp(min=0)
        loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log()
        return loss
    
class BCE_from_logits_focal(nn.modules.Module):
    def __init__(self, gamma):
        super(BCE_from_logits_focal,self).__init__()
        self.gamma = gamma
    def forward(self, input, target):
        #input = input.clamp(min = -1, max = 1)
        max_val = (-input).clamp(min=0)
        loss = input - input * target + max_val + ((-max_val).exp() + (-input - max_val).exp()).log()
        p = input.sigmoid()
        pt = (1-p)*(1-target) + p*target
        return ((1-pt).pow(self.gamma))*loss




from nnunet.training.loss_functions.dice_loss import SoftDiceLossSquared

dice_loss = SoftDiceLossSquared()

import os, sys
#this helper function just redirects the print statments from nnUnet, which are quite verbose :-) 
class HiddenPrints:
    def __enter__(self):
        self._original_stdout = sys.stdout
        sys.stdout = open(os.devnull, 'w')

    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout.close()
        sys.stdout = self._original_stdout


def apply_to_case(model, patient, soft_dice_weighting = 0, do_mirroring=False, patch_size=PATCH_SIZE, axes=['axial'], 
                 size=(192,192)):
    model.eval()
    with torch.no_grad():
      ensemble_logits = []

      patient_data, patient_metadata = load_patient(patient, axis='axial')
      masks = pad_nd_image(patient_data[-1],patch_size)
      full_image, slicer =  pad_nd_image(patient_data[:-1],patch_size, return_slicer=True)
      full_mask =  make_target(masks[np.newaxis,np.newaxis])[0]
      gradient_mask = get_gradient_mask(full_image[np.newaxis,:]).astype(np.uint8)[0]
      for axis in axes:
          
          all_data, patient_metadata = load_patient(patient, axis=axis)
          image_data =  pad_nd_image(all_data[:-1], patch_size)
          zero_mask = (image_data != 0).astype(np.uint8)
          image_data = image_data - np.array(patient_metadata['means'])[:,np.newaxis,np.newaxis,np.newaxis]
          image_data = image_data/np.array(patient_metadata['sds'])[:,np.newaxis,np.newaxis,np.newaxis]

          image_data = image_data * zero_mask
          with HiddenPrints():
            output = unet.predict_3D(torch.from_numpy(image_data),do_mirroring=do_mirroring, patch_size=patch_size,tiled=True)
          
          output = output[2]

          if axis == 'coronal':
              output = np.swapaxes(output, 1,2)

          if axis == 'sagittal':
              output = np.swapaxes(output, 1,3)
          
          ensemble_logits.append(output)



          
      full_logit = np.mean(np.array(ensemble_logits),axis=0)

      return full_mask[slicer], full_logit[slicer]
class Average(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.sum += val
        self.count += n

    @property
    def avg(self):
        return self.sum / self.count
from torch.optim.lr_scheduler import _LRScheduler

class ExponentialLR(_LRScheduler):
    """Exponentially increases the learning rate between two boundaries over a number of
    iterations.
    Arguments:
        optimizer (torch.optim.Optimizer): wrapped optimizer.
        end_lr (float): the final learning rate.
        num_iter (int): the number of iterations over which the test occurs.
        last_epoch (int, optional): the index of last epoch. Default: -1.

        from https://github.com/davidtvs/pytorch-lr-finder/
    """

    def __init__(self, optimizer, end_lr, num_iter, last_epoch=-1):
        self.end_lr = end_lr
        self.num_iter = num_iter
        super(ExponentialLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        curr_iter = self.last_epoch + 1
        r = curr_iter / self.num_iter
        return [base_lr * (self.end_lr / base_lr) ** r for base_lr in self.base_lrs]
import copy

def train_unet(net = None,num_epochs = 50, batches_per_train_epoch = 20, num_val_cases = 1,
                  dice_weighting = DICE_WEIGHTING, lr_max = 0.0001,
                  lr_min = 0.00000001, scheduler = 'one-cycle', lr_find = False, lr_find_steps = 100, evaluate = True):
    
    optimizer = optim.Adam(net.parameters(), lr=lr_max, weight_decay=0.00001)
    loader = tr_gen
    epochs_this_run = 0

    lr_history = []

    loss_history = []

    if lr_find:
      num_epochs = 1
      batches_per_train_epoch = lr_find_steps
      evaluate = False
      initial_model = copy.deepcopy(net.state_dict()) 


    if lr_find:
          for param_group in optimizer.param_groups:
            param_group['lr'] = lr_min
          schedule = ExponentialLR(optimizer, end_lr = lr_max, num_iter = lr_find_steps)
    elif scheduler == 'cosine':
          print('using per-epoch cosine annealing')
          for param_group in optimizer.param_groups:
            param_group['lr'] = lr_max
          schedule = CosineAnnealingLR(optimizer, T_max = batches_per_train_epoch, eta_min=lr_min)
    else:
          print('using one-cycle')
          div_factor = 25
          for param_group in optimizer.param_groups:
            param_group['lr'] = lr_max/25
            schedule = OneCycleLR(optimizer, max_lr = lr_max, epochs = num_epochs, steps_per_epoch = batches_per_train_epoch, div_factor = 25)
    for epoch in tqdm(range(0, num_epochs)):
        
        epochs_this_run = epochs_this_run+1

        if scheduler == 'cosine':
          for param_group in optimizer.param_groups:
            param_group['lr'] = lr_max
          schedule = CosineAnnealingLR(optimizer, T_max = batches_per_train_epoch, eta_min=lr_min)

        train_loss = Average()
        soft_dice_loss = Average()
                
        net.train()
        t = tqdm(range(batches_per_train_epoch),leave=True)

        for iter in  t:
            lr_history.append(schedule.get_lr())
            this_batch = tr_gen.next()
            images = this_batch['data']
            masks = this_batch['seg']

            zero_mask = (masks[:,0]==-1).astype(np.float)

            #images = images*(1-zero_mask)[:,np.newaxis]

            gradient_masks = get_gradient_mask(images)


            masks = torch.from_numpy(make_target(masks))
            images = torch.from_numpy(images)
            gradient_masks = torch.from_numpy(gradient_masks)


            optimizer.zero_grad()

            outputs = net(Variable(images).cuda())

            outputs = outputs[0]*gradient_masks.unsqueeze(1).cuda()

            
            total_loss = 0
            

            masks = (masks).cuda(async=True)
            gradient_mask = (gradient_masks).cuda(async=True) 



            loss = BCE_from_logits()(outputs,masks.cuda())

            soft_dice = 1 + dice_loss(masks.cuda(), outputs.sigmoid())

            total_loss = loss.sum() 


            train_loss.update(total_loss.item(), torch.sum(gradient_mask != 0).item())
            soft_dice_loss.update(soft_dice.item(), 1)


            eps=1e-10
            total_loss = total_loss/(torch.sum(gradient_mask != 0.0) + eps) + soft_dice*dice_weighting
            
            outputs_prob = torch.sigmoid(outputs)

            total_loss.backward()

            optimizer.step()
           # optimizer2.step()
            train_loss.update(total_loss.item(), images.size(0))
            loss_history.append(total_loss.item())

            t.set_description(f'Loss = {train_loss.avg:.4f}, {soft_dice_loss.avg:.4f}')
            schedule.step()
                        
        
        #loss = (criterion(outputs, masks)[0] +(-logit_flip).exp()[0]).data.cpu().numpy()
        
        all_dices = []

        if evaluate:
          val_loss = Average()

          val_dice_loss = Average()
          
          
          val_dice = [Average() for x in TARGET_LABEL_SETS]
          
          net.eval()
          
          
          result_summary = list()

          all_dices = []
          for case in tqdm(range(num_val_cases), leave=False): 
              patient = val[case]

              patient_data, patient_metadata = load_patient(patient, axis='axial')              
              gradient_mask = get_gradient_mask(patient_data[:-1][np.newaxis,:]).astype(np.uint8)[0]
              mask, logit = apply_to_case(net,patient=patient, 
                                        axes=['axial','sagittal','coronal'])
              
              logit = logit * np.expand_dims(gradient_mask,0)
      
              TP = np.sum(np.logical_and(logit>0 , mask>0), axis=(1,2,3))
              FP = np.sum(np.logical_and(logit>0 , mask==0), axis=(1,2,3))
              FN = np.sum(np.logical_and(logit<0 , mask>0), axis=(1,2,3))
              TN = np.sum(np.logical_and(logit<0 , mask==0), axis=(1,2,3))
      

              loss = BCE_from_logits()(torch.Tensor(logit), torch.Tensor(mask)).numpy()
              soft_dice = 1 + dice_loss(torch.Tensor(mask).unsqueeze(0), torch.Tensor(logit).unsqueeze(0).sigmoid()).numpy()
            
              

              
              val_loss.update(np.sum(loss), np.sum(gradient_mask != 0.0))
              
              val_dice_loss.update(soft_dice, 1)

              epsilon = 0.000001
              dices = (2*TP+epsilon)/(2*TP+FP+FN+epsilon)


              all_dices.append(dices)
        
          print(f'Epoch = {epoch},  Dices ' , [f'{x}:{y:.2f}' for x, y in zip(TARGET_LABEL_NAMES, np.mean(all_dices, 0))], 
                f'BCE loss = {val_loss.avg:.2f}',
                f'Dice loss = {val_dice_loss.avg:.2f}')  
    if lr_find:
      print('restoring model')
      net.load_state_dict(initial_model)      
    return net, all_dices, lr_history, loss_history
num_val_cases
unet = Generic_UNet(input_channels=4, base_num_features=12, num_classes=3, num_pool=4, num_conv_per_stage=2,
                 feat_map_mul_on_downscale=2, conv_op=nn.Conv3d,
                 norm_op=nn.InstanceNorm3d, norm_op_kwargs=None,
                 dropout_op=nn.Dropout3d, dropout_op_kwargs=None,
                 nonlin=nn.LeakyReLU,deep_supervision=False,final_nonlin = nn.Identity(), seg_output_use_bias=True)

pytorch_total_params = sum(np.fromiter((p.numel() for p in unet.parameters()), int))
print(pytorch_total_params)

unet = unet.cuda()

pi = np.array(label_prevalences)

new_bias = torch.from_numpy(-np.log(((1-pi)/pi)))


unet.seg_outputs[-1].bias = torch.nn.Parameter(new_bias.float().cuda())



unet.seg_outputs[-1].bias
tr_gen.restart()

net, dices, lr_history, loss_history = train_unet(unet, NUM_EPOCHS,batches_per_train_epoch=1000, dice_weighting = DICE_WEIGHTING, lr_max=1e-1, lr_min=1e-8, 
                            lr_find = True, lr_find_steps = 500)
unet.seg_outputs[-1].bias
from scipy.ndimage.filters import gaussian_filter1d

ysmoothed = gaussian_filter1d(loss_history, sigma=8)



plt.plot(lr_history, ysmoothed)
plt.xscale('log')
plt.ylim(np.min(ysmoothed),1.2*ysmoothed[0])

NUM_EPOCHS=4


net, dices, lr_history, loss_history = train_unet(unet, NUM_EPOCHS,batches_per_train_epoch=BATCHES_PER_EPOCH, dice_weighting = DICE_WEIGHTING, lr_max=1e-3,scheduler = 'one_cycle', num_val_cases = num_val_cases)
plt.plot(lr_history)
from scipy.ndimage.filters import gaussian_filter1d


plt.plot(gaussian_filter1d(loss_history, sigma=100))
tr_gen.restart()

net, dices, lr_history, loss_history = train_unet(unet, NUM_EPOCHS,batches_per_train_epoch=1000, dice_weighting = DICE_WEIGHTING, lr_max=1e-2, lr_min=1e-8, 
                            lr_find = True, lr_find_steps = 500)
ysmoothed = gaussian_filter1d(loss_history, sigma=8)



plt.plot(lr_history, ysmoothed)
plt.xscale('log')
plt.ylim(np.min(ysmoothed),1.2*ysmoothed[0])
tr_gen.restart()

NUM_EPOCHS=4


net, dices, lr_history, loss_history = train_unet(unet, NUM_EPOCHS,batches_per_train_epoch=BATCHES_PER_EPOCH, dice_weighting = DICE_WEIGHTING, lr_max=1e-4,scheduler = 'one_cycle', num_val_cases = num_val_cases)
from scipy.ndimage.filters import gaussian_filter1d


plt.plot(gaussian_filter1d(loss_history, sigma=100))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09639433097123802, 'fraction_numerical': 0.016882034180908712, 'mean_word_length': 3.175344452501813, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 19, 'https://': 6, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '282757', 'license': 'Apache-2.0', 'n_tokens_mistral': 9769, 'n_tokens_neox': 9166, 'n_words': 2366, 'path': '_notebooks/2020-04-23-BraTS_3D__with_dice_loss_(nnUnet).ipynb', 'repo_name': 'RIMcKinley/DeepSCANblog', 'size': '28788'}