training loop for segmenter, brain mri dataset

train_segmenter.py

def run_experiment(model_name, model, optimizer, criterion, train_loader, val_loader, device='cuda', num_epochs=50, clear_mem=True):
    
    #######################
    #Train model          #
    #######################
    print('Model sent to ' + str(device))
    model.to(device)
    losses = []
    train_scores = [] # hold IoU scores
    iters = 0
    for epoch in range(num_epochs):
        if epoch % 10 == 0:
            print(f'Epoch {epoch+1}/{num_epochs}')
        
        for i,batch in enumerate(train_loader):
            img = batch[0].to(device)
            msk = batch[1].to(device)

            optimizer.zero_grad()
            output = model(img)
            loss = criterion(output, msk)
            loss.backward()
            optimizer.step()

            losses.append(loss.item())
            train_scores.append(iou_pytorch(output.detach(), msk))

            iters += 1

            #if iters % 500 == 0:
                #print(f'Loss: [{loss}]')
                
    #for i in range(len(train_scores)):
     #   train_scores[i] = train_scores[i].mean()
        
    #######################
    #Validate model       #
    #######################
    
    model.eval()
    val_losses = []
    val_scores = []

    for i,batch in enumerate(val_loader):
            img = batch[0].to(device)
            msk = batch[1].to(device)

            output = model(img)
            loss = criterion(output, msk)
            val_scores.append(iou_pytorch(output.detach(), msk))
            val_losses.append(loss.item())
            
    results = {
        'model_name': model_name,
        'train_losses': losses,
        'train_scores': train_scores,
        'val_losses': val_losses,
        'val_scores': val_scores
    }
    
    if clear_mem:
        del model
        del optimizer
        del criterion
        torch.cuda.empty_cache()
        
    return results