train_correct.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torchvision import datasets, transforms
from torch.optim import Optimizer
from torch.utils import data
import pretrainedmodels
import numpy as np
import os
import cv2
from skimage.io import imread
from torch.utils.data.sampler import WeightedRandomSampler, BatchSampler
from albumentations import (
    HorizontalFlip, VerticalFlip, IAAPerspective, ShiftScaleRotate, CLAHE, RandomRotate90, Normalize, RandomGamma, RandomBrightnessContrast, HueSaturationValue, CLAHE, ChannelShuffle, 
    Transpose, ShiftScaleRotate, Blur, OpticalDistortion, GridDistortion, HueSaturationValue,
    IAAAdditiveGaussianNoise, GaussNoise, MotionBlur, MedianBlur, IAAPiecewiseAffine,
    IAASharpen, IAAEmboss, RandomContrast, RandomBrightness, Flip, OneOf, Compose, PadIfNeeded, RandomCrop, Resize
)
import pretrainedmodels.utils as utils
from sklearn.metrics import roc_auc_score

def train(model, train_loader, optimizer, epoch, log_interval, loss_f, samples_per_epoch, device, cycling_optimizer=False):
    """Trains the model using the provided optimizer and loss function.
    Shows output each log_interval iterations 
    Args:
        model: Pytorch model to train.
        train_loader: Data loader.
        optimizer: pytroch optimizer.
        epoch: Current epoch.
        log_interval: Show model training progress each log_interval steps.
        loss_f: Loss function to optimize.
        samples_per_epoch: Number of samples per epoch to scale loss.
        device: pytorch device
        cycling_optimizer: Indicates of optimizer is cycling.
    """
    model.train()
    total_losses = []
    losses =[]
    for batch_idx, (x, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(x.to(device, dtype=torch.float))
        loss = loss_f(output, target.to(device, dtype=torch.float))
        losses.append(loss.item())
        loss.backward()
        nn.utils.clip_grad_norm(model.parameters(), 4)
        if cycling_optimizer:
            optimizer.batch_step()
        else:
            optimizer.step()
        if batch_idx % log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.3f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(x), samples_per_epoch,
                100. * batch_idx * len(x) / samples_per_epoch, np.mean(losses)))
            total_losses.append(np.mean(losses))
            losses = []
    train_loss_mean = np.mean(total_losses)
    print('Mean train loss on epoch {} : {}'.format(epoch, train_loss_mean))
    return train_loss_mean
  
  def test(model, test_loader, loss_f, epoch, device):
    """Test the model with validation data.
    Args:
        model: Pytorch model to test data with.
        test_loader: Data loader.
        loss_f: Loss function.
        epoch: Current epoch.
        device: pytorch device        
    """
    model.eval()
    test_loss = 0
    predictions=[]
    targets=[]
    test_loss=[]
    with torch.no_grad():
        for x, target in test_loader:
            output = model(x.to(device, dtype=torch.float))
            test_loss.append(loss_f(output, target.to(device, dtype=torch.float)).item())
            predictions.append(output.cpu())
            targets.append(target.cpu())
    predictions = np.vstack(predictions)
    targets = np.vstack(targets)
    score = roc_auc_score(targets, predictions)
    test_loss  = np.mean(test_loss)
    print('\nTest set: Average loss: {:.6f}, roc auc: {:.4f}\n'.format(test_loss, score))
    return test_loss, score	
  
  lass Net(nn.Module):
    """Build the nn network based on pretrained resnet models.
    Args:
        base_model: resnet34\resnet50\etc from pretrained models
        n_features: n features from last pooling layer       
    """
    def __init__(self, base_model, n_features):
        super(Net, self).__init__()
        self.layer0 = nn.Sequential(*list(base_model.children())[:4])
        self.layer1 = nn.Sequential(*list(base_model.layer1))
        self.layer2 = nn.Sequential(*list(base_model.layer2))
        self.layer3 = nn.Sequential(*list(base_model.layer3))
        self.layer4 = nn.Sequential(*list(base_model.layer4))
        self.dense1 = nn.Sequential(nn.Linear(n_features, 128))
        self.dense2 = nn.Sequential(nn.Linear(128, 64))
        self.classif = nn.Sequential(nn.Linear(64, 1))
    def forward(self, x):
        x = self.features(x)
        x = F.avg_pool2d(x, 7)
        x = x.view(x.size(0), -1)
        x = self.dense1(x)
        x = self.dense2(x)
        x = self.classif(x)
        x = torch.sigmoid(x)
        return x
    def features(self, x):
        x = self.layer0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        return x