classify_moa.py

#!/usr/bin/env python3

import train_funcs

cell_lines = [
    "MDA-157",
    "MDA-231",
    "MCF7",
    "KPL4",
    "T47D",
    "HCC1954",
    "HCC1569",
    "SKBR3"
]

for cell_line in cell_lines:
    train_funcs.train_on_cell_line(cell_line)

train_funcs.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import copy
import os
import json


def train_on_cell_line(cell_line):
    # Data augmentation and normalization for training
    # Just normalization for validation
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomSizedCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Scale(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ]),
    }

    data_dir = '/home/scott/chopped/nncell_data_300_{}/'.format(cell_line)
    dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
             for x in ['train', 'test']}
    dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=16,
                                                   shuffle=True, num_workers=8)
                    for x in ['train', 'test']}
    dset_sizes = {x: len(dsets[x]) for x in ['train', 'test']}
    dset_classes = dsets['train'].classes

    use_gpu = torch.cuda.is_available()

    print(" --- Training on {} cell_line --- ".format(cell_line))

    def plot_history(history, plot_path):
        """plot training history"""
        plt.figure(figsize=[8, 9])
        plt.subplot(211)
        plt.grid(linestyle="--", alpha=0.5)
        plt.plot(history["test_acc"], label="test_acc")
        plt.plot(history["train_acc"], label="train_acc")
        plt.legend()
        plt.title(cell_line)

        plt.subplot(212)
        plt.grid(linestyle="--", alpha=0.5)
        plt.plot(history["test_loss"], label="test_loss")
        plt.plot(history["train_loss"], label="train_loss")
        plt.legend()

        plt.tight_layout()
        plt.savefig(plot_path)
        plt.close()



    def train_model(model, criterion, optimizer, lr_scheduler,
                    plot_path, num_epochs=100):

        since = time.time()

        best_model = model
        best_acc = 0.0

        if use_gpu is False:
            raise RuntimeError("No GPU device found, aborting")
        else:
            print("GPU device found")

        history = {
            "test_acc"   : [],
            "test_loss"   : [],
            "train_acc"  : [],
            "train_loss" : []
        }


        for epoch in range(num_epochs):
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            print('-' * 10)

            # Each epoch has a training and validation phase
            for phase in ['train', 'test']:
                if phase == 'train':
                    optimizer = lr_scheduler(optimizer, epoch)
                    model.train(True)  # Set model to training mode
                else:
                    model.train(False)  # Set model to evaluate mode

                running_loss = 0.0
                running_corrects = 0

                # Iterate over data.
                for data in dset_loaders[phase]:
                    # get the inputs
                    inputs, labels = data

                    # wrap them in Variable
                    if use_gpu:
                        inputs, labels = Variable(inputs.cuda()), \
                            Variable(labels.cuda())
                    else:
                        inputs, labels = Variable(inputs), Variable(labels)

                    # zero the parameter gradients
                    optimizer.zero_grad()

                    # forward
                    outputs = model(inputs)
                    _, preds = torch.max(outputs.data, 1)
                    loss = criterion(outputs, labels)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                    # statistics
                    running_loss += loss.data[0]
                    running_corrects += torch.sum(preds == labels.data)

                epoch_loss = running_loss / dset_sizes[phase]
                epoch_acc = running_corrects / dset_sizes[phase]
                history["{}_acc".format(phase)].append(epoch_acc)
                history["{}_loss".format(phase)].append(epoch_loss)

                plot_history(history, plot_path)

                print('{} Loss: {:.4f} Acc: {:.4f}'.format(
                    phase, epoch_loss, epoch_acc))

                # deep copy the model
                if phase == 'test' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    best_model = copy.deepcopy(model)

            print()

        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(
            time_elapsed // 60, time_elapsed % 60))
        print('Best test Acc: {:4f}'.format(best_acc))
        return best_model, history




    def exp_lr_scheduler(optimizer, epoch, init_lr=0.005, lr_decay_epoch=20):
        """Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs."""
        lr = init_lr * (0.1**(epoch // lr_decay_epoch))

        if epoch % lr_decay_epoch == 0:
            print('LR is set to {}'.format(lr))

        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        return optimizer



    model_ft = models.resnet18(pretrained=False)
    num_ftrs = model_ft.fc.in_features
    model_ft.fc = nn.Linear(num_ftrs, 8)

    if use_gpu:
        model_ft = model_ft.cuda()

    criterion = nn.CrossEntropyLoss()

    # Observe that all parameters are being optimized
    optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

    plot_path = "../plots/{}_training.pdf".format(cell_line)

    model_ft, history = train_model(
        model_ft, criterion, optimizer_ft,
        exp_lr_scheduler, num_epochs=250, plot_path=plot_path
    )

    # save the model
    model_path = "../models/{}_trained_model.pynn".format(cell_line)
    torch.save(model_ft.state_dict(), model_path)

    # write history dict as a json file
    with open("../models/history_{}.json".format(cell_line), "w") as f:
        json.dump(history, f, indent=4)