train.py

#!/usr/bin/env python

"""
Training of models
"""

###########
# Imports #
###########

import time
import torch


#############
# Functions #
#############

def train_epoch(model, loader, criterion, optimizer):
	'''Trains a model for one epoch.'''
	model.train()
	device = next(model.parameters()).device
	losses = []

	for inputs, targets in loader:
		inputs = inputs.to(device)
		targets = targets.to(device)
		outputs = model(inputs)

		loss = criterion(outputs, targets)
		losses.append(loss.item())

		optimizer.zero_grad()
		loss.backward()
		optimizer.step()

	return losses


########
# Main #
########

if __name__ == '__main__':
	# Imports
	import argparse
	import csv
	import numpy as np
	import os
	import random
	import sys
	import via as VIA

	from torch.utils.data import DataLoader
	from torch.optim import Adam, SGD

	from dataset import VIADataset, ColorJitter, RandomFilter, RandomTranspose, Scale, ToTensor
	from models import UNet, SegNet, MultiTaskUNet, MultiTaskSegNet
	from criterions import DiceLoss, MultiTaskLoss, TP, TN, FP, FN

	# Arguments
	parser = argparse.ArgumentParser(description='Train a model')
	parser.add_argument('-d', '--destination', default='../products/models/', help='destination of the network file(s)')
	parser.add_argument('-e', '--epochs', type=int, default=100, help='number of epochs')
	parser.add_argument('-f', '--fold', type=int, default=0, help='the fold')
	parser.add_argument('-i', '--input', default='../products/json/california.json', help='input VIA file')
	parser.add_argument('-k', type=int, default=5, help='the number of folds')
	parser.add_argument('-m', '--model', default='unet', choices=['unet', 'segnet'], help='network schema')
	parser.add_argument('-multitask', default=False, action='store_true', help='multi-task network')
	parser.add_argument('-n', '--name', default=None, help='name of the network')
	parser.add_argument('-o', '--output', default=None, help='standard output file')
	parser.add_argument('-p', '--path', default='../resources/california/', help='path to resources')
	parser.add_argument('-r', '--resume', type=int, default=1, help='epoch at which to resume')
	parser.add_argument('-s', '--stat', default='../products/csv/statistics.csv', help='convergence statistics file')
	parser.add_argument('-scale', type=int, default=1, help='scale of the images')
	parser.add_argument('-batch', type=int, default=5, help='batch size')
	parser.add_argument('-optim', default='adam', choices=['adam', 'sgd'], help='optimizer')
	parser.add_argument('-lrate', type=float, default=1e-3, help='learning rate')
	parser.add_argument('-wdecay', type=float, default=1e-4, help='weight decay')
	parser.add_argument('-momentum', type=float, default=0.9, help='momentum of SGD')
	parser.add_argument('-special', default=False, action='store_true', help='special mode')
	args = parser.parse_args()

	# Output file
	if args.output is not None:
		if os.path.dirname(args.output):
			os.makedirs(os.path.dirname(args.output), exist_ok=True)
		sys.stdout = open(args.output, 'a')

	print('-' * 10)

	# Datasets
	via = VIA.load(args.input)

	keys = sorted(list(via.keys()))

	random.seed(0) # reproductability
	random.shuffle(keys)

	if (args.k > 0):
		train_via = {key: via[key] for i, key in enumerate(keys) if (i % args.k) != args.fold}
	else:
		train_via = via

	if args.special:
		trainset = VIADataset(train_via, args.path, shuffle=True, size=None)
		trainset = RandomTranspose(trainset)
	else:
		trainset = VIADataset(train_via, args.path, shuffle=True, alt=1)
		trainset = Scale(trainset, args.scale) if args.scale > 1 else trainset
		trainset = RandomTranspose(RandomFilter(ColorJitter(trainset)))

	trainset = ToTensor(trainset)

	print('Training size = {}'.format(len(trainset)))

	# Dataloaders
	trainloader = DataLoader(trainset, batch_size=args.batch, pin_memory=True)

	# Model
	if args.model == 'unet':
		if args.multitask:
			model = MultiTaskUNet(3, 1, R=5)
		else:
			model = UNet(3, 1)
	elif args.model == 'segnet':
		if args.multitask:
			model = MultiTaskSegNet(3, 1, R=5)
		else:
			model = SegNet(3, 1)
	else:
		raise ValueError('unknown model {}'.format(args.model))

	if torch.cuda.is_available():
		print('CUDA available -> Transfering to CUDA')
		device = torch.device('cuda')
	else:
		device = torch.device('cpu')
		print('CUDA unavailable')

	model = model.to(device)

	os.makedirs(args.destination, exist_ok=True)
	basename = os.path.join(
		args.destination,
		args.model if args.name is None else args.name
	)

	if args.resume > 1:
		modelname = '{}_{:03d}.pth'.format(basename, args.resume - 1)

		if os.path.exists(modelname):
			print('Resuming from {}'.format(modelname))
			model.load_state_dict(torch.load(modelname, map_location=device))
		else:
			args.resume = 1
	else:
		args.resume = 1

	# Convergence statistics
	if os.path.dirname(args.stat):
		os.makedirs(os.path.dirname(args.stat), exist_ok=True)

	if not os.path.exists(args.stat):
		with open(args.stat, 'w', newline='') as f:
			csv.writer(f).writerow([
				'model',
				'epoch',
				'train_loss_mean',
				'train_loss_std',
				'train_loss_first',
				'train_loss_second',
				'train_loss_third'
			])

	# Criterion
	if args.multitask:
		train_criterion = MultiTaskLoss(smooth=1., R=5)
	else:
		train_criterion = DiceLoss(smooth=1.)

	# Optimizer
	if args.optim == 'adam':
		optimizer = Adam(
			model.parameters(),
			lr=args.lrate,
			weight_decay=args.wdecay
		)
	elif args.optim == 'sgd':
		optimizer = SGD(
			model.parameters(),
			lr=args.lrate,
			weight_decay=args.wdecay,
			momentum=args.momentum
		)
	else:
		raise ValueError('unknown optimizer {}'.format(args.optim))

	# Training
	epochs = range(args.resume, args.resume + args.epochs)

	for epoch in epochs:
		if args.output is not None:
			sys.stdout.close()
			sys.stdout = open(args.output, 'a')

		print('-' * 10)
		print('Epoch {}'.format(epoch))

		start = time.time()

		## Training set
		train_losses = train_epoch(model, trainloader, train_criterion, optimizer)

		elapsed = time.time() - start

		print('{:.0f}m{:.0f}s elapsed'.format(elapsed // 60, elapsed % 60))

		## Statistics
		train_losses = np.array(train_losses)

		train_mean = train_losses.mean()

		print('Training loss = {}'.format(train_mean))

		with open(args.stat, 'a', newline='') as f:
			csv.writer(f).writerow([
				args.name,
				epoch,
				np.mean(train_losses),
				np.std(train_losses),
				np.quantile(train_losses, 0.25),
				np.quantile(train_losses, 0.5),
				np.quantile(train_losses, 0.75)
			])

		## Saving last
		if epoch == epochs[-1]:
			modelname = '{}_{:03d}.pth'.format(basename, epoch)
			print('Saving {}'.format(modelname))
			torch.save(model.state_dict(), modelname)