Commit 05ed6c03 authored by Anjith GEORGE's avatar Anjith GEORGE Committed by Anjith GEORGE

Trainers and extractores for FASNet

parent a75bdc4e
import numpy as np
import torch
from torch.autograd import Variable
import torchvision.transforms as transforms
from bob.learn.pytorch.architectures import FASNet
from bob.bio.base.extractor import Extractor
import logging
logger = logging.getLogger("bob.learn.pytorch")
class FASNetExtractor(Extractor):
""" The class implementing the FASNet score computation.
Attributes
----------
network: :py:class:`torch.nn.Module`
The network architecture
transforms: :py:mod:`torchvision.transforms`
The transform from numpy.array to torch.Tensor
"""
def __init__(self, transforms = transforms.Compose([transforms.ToPILImage(),transforms.Resize(224, interpolation=2),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])]), model_file=None):
""" Init method
Parameters
----------
num_channels_used: int
The number of channels to be used by the network. This could be
different from the number of channels present in the input image. For instance,
when used together with 'ChannelSelect' transform. The value of `num_channels_used`
should be the number of channels eventually used by the network (i.e., output of transform).
model_file: str
The path of the trained PAD network to load
transforms: :py:mod:`torchvision.transforms`
tranform to be applied on the image
"""
Extractor.__init__(self, skip_extractor_training=True)
# model
self.transforms = transforms
self.network = FASNet()
#self.network=self.network.to(device)
if model_file is None:
# do nothing (used mainly for unit testing)
logger.debug("No pretrained file provided")
pass
else:
# With the new training
logger.debug('Starting to load the pretrained PAD model')
cp = torch.load(model_file)
if 'state_dict' in cp:
self.network.load_state_dict(cp['state_dict'])
logger.debug('Loaded the pretrained PAD model')
self.network.eval()
def __call__(self, image):
""" Extract features from an image
Parameters
----------
image : 3D :py:class:`numpy.ndarray` (floats)
The multi-channel image to extract the score from. Its size must be num_channelsx128x128;
Note: the value of `num_channels` is the number of channels as obtained from the preprocessed
data. The actual number of channels used may vary, for instance
if `ChannelSelect` transform is used, the number of channels used would change.
Returns
-------
output : float
The extracted feature is a scalar values ~1 for bonafide and ~0 for PAs
"""
input_image = np.rollaxis(np.rollaxis(image, 2),2) # changes to 128x128xnum_channels
input_image = self.transforms(input_image)
input_image = input_image.unsqueeze(0)
output = self.network.forward(Variable(input_image))
output = output.data.numpy().flatten()
# output is a scalar score
return output
\ No newline at end of file
......@@ -3,6 +3,7 @@ from .LightCNN29 import LightCNN29Extractor
from .LightCNN29v2 import LightCNN29v2Extractor
from .MCCNN import MCCNNExtractor
from .MCCNNv2 import MCCNNv2Extractor
from .FASNet import FASNetExtractor
__all__ = [_ for _ in dir() if not _.startswith('_')]
#!/usr/bin/env python
# encoding: utf-8
""" Train a FASNet for face PAD
Usage:
%(prog)s <configuration>
[--model=<string>] [--batch-size=<int>] [--num-workers=<int>][--epochs=<int>]
[--learning-rate=<float>][--do-crossvalidation][--seed=<int>]
[--output-dir=<path>] [--use-gpu] [--verbose ...]
Arguments:
<configuration> A configuration file, defining the dataset and the network
Options:
-h, --help Shows this help message and exits
--model=<string> Filename of the model to load (if any).
--batch-size=<int> Batch size [default: 64]
--num-workers=<int> Number subprocesses to use for data loading [default: 0]
--epochs=<int> Number of training epochs [default: 20]
--learning-rate=<float> Learning rate [default: 0.01]
--do-crossvalidation Whether to perform cross validation [default: False]
-S, --seed=<int> The random seed [default: 3]
-o, --output-dir=<path> Dir to save stuff [default: training]
-g, --use-gpu Use the GPU
-v, --verbose Increase the verbosity (may appear multiple times).
Note that arguments provided directly by command-line will override the ones in the configuration file.
Example:
To run the training process
$ %(prog)s config.py
See '%(prog)s --help' for more information.
"""
import os, sys
import pkg_resources
import torch
import numpy
from docopt import docopt
import bob.core
logger = bob.core.log.setup("bob.learn.pytorch")
from bob.extension.config import load
from bob.learn.pytorch.trainers import FASNetTrainer
from bob.learn.pytorch.utils import get_parameter
version = pkg_resources.require('bob.learn.pytorch')[0].version
def main(user_input=None):
# Parse the command-line arguments
if user_input is not None:
arguments = user_input
else:
arguments = sys.argv[1:]
prog = os.path.basename(sys.argv[0])
completions = dict(prog=prog, version=version,)
args = docopt(__doc__ % completions,argv=arguments,version='Train a FASNet (%s)' % version,)
# load configuration file
configuration = load([os.path.join(args['<configuration>'])])
# get the pre-trained model file, if any
model = args['--model']
if hasattr(configuration, 'model'):
model = configuration.model
# get various parameters, either from config file or command-line
batch_size = get_parameter(args, configuration, 'batch_size', 64)
num_workers = get_parameter(args, configuration, 'num_workers', 0)
epochs = get_parameter(args, configuration, 'epochs', 20)
learning_rate = get_parameter(args, configuration, 'learning_rate', 0.01)
seed = get_parameter(args, configuration, 'seed', 3)
output_dir = get_parameter(args, configuration, 'output_dir', 'training')
use_gpu = get_parameter(args, configuration, 'use_gpu', False)
verbosity_level = get_parameter(args, configuration, 'verbose', 0)
do_crossvalidation = get_parameter(args, configuration, 'do_crossvalidation', False)
bob.core.log.set_verbosity_level(logger, verbosity_level)
bob.io.base.create_directories_safe(output_dir)
# print parameters
logger.debug("Model file = {}".format(model))
logger.debug("Batch size = {}".format(batch_size))
logger.debug("Num workers = {}".format(num_workers))
logger.debug("Epochs = {}".format(epochs))
logger.debug("Learning rate = {}".format(learning_rate))
logger.debug("Seed = {}".format(seed))
logger.debug("Output directory = {}".format(output_dir))
logger.debug("Use GPU = {}".format(use_gpu))
logger.debug("Perform cross validation = {}".format(do_crossvalidation))
# process on the arguments / options
torch.manual_seed(seed)
if use_gpu:
torch.cuda.manual_seed_all(seed)
if torch.cuda.is_available() and not use_gpu:
logger.warn("You have a CUDA device, so you should probably run with --use-gpu")
# get data
if hasattr(configuration, 'dataset'):
dataloader={}
if not do_crossvalidation:
logger.info("There are {} training samples".format(len(configuration.dataset['train'])))
dataloader['train'] = torch.utils.data.DataLoader(configuration.dataset['train'], batch_size=batch_size, num_workers=num_workers, shuffle=True)
else:
dataloader['train'] = torch.utils.data.DataLoader(configuration.dataset['train'], batch_size=batch_size, num_workers=num_workers, shuffle=True)
dataloader['val'] = torch.utils.data.DataLoader(configuration.dataset['val'], batch_size=batch_size, num_workers=num_workers, shuffle=True)
logger.info("There are {} training samples".format(len(configuration.dataset['train'])))
logger.info("There are {} validation samples".format(len(configuration.dataset['val'])))
else:
logger.error("Please provide a dataset in your configuration file !")
sys.exit()
# train the network
if hasattr(configuration, 'network'):
trainer = FASNetTrainer(configuration.network, batch_size=batch_size, use_gpu=use_gpu, verbosity_level=verbosity_level,tf_logdir=output_dir+'/tf_logs',do_crossvalidation=do_crossvalidation)
trainer.train(dataloader, n_epochs=epochs, learning_rate=learning_rate, output_dir=output_dir, model=model)
else:
logger.error("Please provide a network in your configuration file !")
sys.exit()
#!/usr/bin/env python
# encoding: utf-8
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from bob.learn.pytorch.utils import comp_bce_loss_weights
from .tflog import Logger
import bob.core
logger = bob.core.log.setup("bob.learn.pytorch")
import time
import os
import copy
class FASNetTrainer(object):
"""
Class to train the MCCNN
Attributes
----------
network: :py:class:`torch.nn.Module`
The network to train
batch_size: int
The size of your minibatch
use_gpu: bool
If you would like to use the gpu
verbosity_level: int
The level of verbosity output to stdout
"""
def __init__(self, network, batch_size=64, use_gpu=False, verbosity_level=2, tf_logdir='tf_logs',do_crossvalidation=False):
""" Init function . The layers to be adapted in the network is selected and the gradients are set to `True`
for the layers which needs to be adapted.
Parameters
----------
network: :py:class:`torch.nn.Module`
The network to train
batch_size: int
The size of your minibatch
use_gpu: bool
If you would like to use the gpu
adapted_layers: str
The blocks in the CNN to adapt; only the ones listed are adapted in the training. The layers are separated by '-' in the
string, for example 'conv1-block1-group1-ffc'. The fully connected layer in the output part are adapted always.
adapt_reference_channel: bool
If this value is `True` then 'ch_0' (which is the reference channel- usually, grayscale image) is also adapted. Otherwise the reference channel
is not adapted, so that it can be used for Face recognition as well, default: `False`.
verbosity_level: int
The level of verbosity output to stdout
"""
self.network = network
self.batch_size = batch_size
self.use_gpu = use_gpu
self.criterion = nn.BCELoss()
self.do_crossvalidation=do_crossvalidation
if self.do_crossvalidation:
phases=['train','val']
else:
phases=['train']
self.phases=phases
if self.use_gpu:
self.network.cuda()
bob.core.log.set_verbosity_level(logger, verbosity_level)
self.tf_logger = Logger(tf_logdir)
# Setting the gradients to true for the layers which needs to be adapted
for name, param in self.network.named_parameters():
param.requires_grad = False
if not 'enc' in name:
param.requires_grad = True
def load_model(self, model_filename):
"""Loads an existing model
Parameters
----------
model_file: str
The filename of the model to load
Returns
-------
start_epoch: int
The epoch to start with
start_iteration: int
The iteration to start with
losses: list(float)
The list of losses from previous training
"""
cp = torch.load(model_filename)
self.network.load_state_dict(cp['state_dict'])
start_epoch = cp['epoch']
start_iter = cp['iteration']
losses = cp['loss']
return start_epoch, start_iter, losses
def save_model(self, output_dir, epoch=0, iteration=0, losses=None):
"""Save the trained network
Parameters
----------
output_dir: str
The directory to write the models to
epoch: int
the current epoch
iteration: int
the current (last) iteration
losses: list(float)
The list of losses since the beginning of training
"""
saved_filename = 'model_{}_{}.pth'.format(epoch, iteration)
saved_path = os.path.join(output_dir, saved_filename)
logger.info('Saving model to {}'.format(saved_path))
cp = {'epoch': epoch,
'iteration': iteration,
'loss': losses,
'state_dict': self.network.cpu().state_dict()
}
torch.save(cp, saved_path)
# moved the model back to GPU if needed
if self.use_gpu :
self.network.cuda()
def train(self, dataloader, n_epochs=25, learning_rate=1e-4, output_dir='out', model=None):
"""Performs the training.
Parameters
----------
dataloader: :py:class:`torch.utils.data.DataLoader`
The dataloader for your data
n_epochs: int
The number of epochs you would like to train for
learning_rate: float
The learning rate for Adam optimizer.
output_dir: str
The directory where you would like to save models
model: str
The path to a pretrained model file to start training from; this is the PAD model; not the LightCNN model
"""
# if model exists, load it
if model is not None:
start_epoch, start_iter, losses = self.load_model(model)
logger.info('Starting training at epoch {}, iteration {} - last loss value is {}'.format(start_epoch, start_iter, losses[-1]))
else:
start_epoch = 0
start_iter = 0
losses = []
logger.info('Starting training from scratch')
for name, param in self.network.named_parameters():
if param.requires_grad == True:
logger.info('Layer to be adapted from grad check : {}'.format(name))
# setup optimizer
optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.network.parameters()),lr = learning_rate )
self.network.train(True)
best_model_wts = copy.deepcopy(self.network.state_dict())
best_loss = float("inf")
# let's go
for epoch in range(start_epoch, n_epochs):
# in the epoch
train_loss_history=[]
val_loss_history = []
for phase in self.phases:
if phase == 'train':
self.network.train() # Set model to training mode
else:
self.network.eval() # Set model to evaluate mode
for i, data in enumerate(dataloader[phase], 0):
if i >= start_iter:
start = time.time()
img, labels = data
labels=labels.float().unsqueeze(1)
weights=comp_bce_loss_weights(labels)
batch_size = len(img)
if self.use_gpu:
img = img.cuda()
labels = labels.cuda()
weights = weights.cuda()
imagesv = Variable(img)
labelsv = Variable(labels)
# weights for samples, should help with data imbalance
self.criterion.weight = weights
optimizer.zero_grad()
with torch.set_grad_enabled(phase == 'train'):
output= self.network(imagesv)
loss = self.criterion(output, labelsv)
if phase == 'train':
loss.backward()
optimizer.step()
train_loss_history.append(loss.item())
else:
val_loss_history.append(loss.item())
end = time.time()
logger.info("[{}/{}][{}/{}] => Loss = {} (time spent: {}), Phase {}".format(epoch, n_epochs, i, len(dataloader[phase]), loss.item(), (end-start),phase))
losses.append(loss.item())
epoch_train_loss=np.mean(train_loss_history)
logger.info("Train Loss : {} epoch : {}".format(epoch_train_loss,epoch))
if self.do_crossvalidation:
epoch_val_loss=np.mean(val_loss_history)
logger.info("Val Loss : {} epoch : {}".format(epoch_val_loss,epoch))
if phase == 'val' and epoch_val_loss < best_loss:
logger.debug("New val loss : {} is better than old: {}, copying over the new weights".format(epoch_val_loss,best_loss))
best_loss = epoch_val_loss
best_model_wts = copy.deepcopy(self.network.state_dict())
######################################## <Logging> ###################################
if self.do_crossvalidation:
info = {'train_loss':epoch_train_loss,'val_loss':epoch_val_loss}
else:
info = {'train_loss':epoch_train_loss}
# scalar logs
for tag, value in info.items():
self.tf_logger.scalar_summary(tag, value, epoch+1)
# Log values and gradients of the parameters (histogram summary)
for tag, value in self.network.named_parameters():
tag = tag.replace('.', '/')
try:
self.tf_logger.histo_summary(tag, value.data.cpu().numpy(), epoch+1)
self.tf_logger.histo_summary(tag+'/grad', value.grad.data.cpu().numpy(), epoch+1)
except:
pass
# Log images
# logimg=img.view(-1,img.size()[1]*224, 128)[:10].cpu().numpy()
# info = { 'images': logimg}
# for tag, images in info.items():
# self.tf_logger.image_summary(tag, images, epoch+1)
######################################## </Logging> ###################################
# do stuff - like saving models
logger.info("EPOCH {} DONE".format(epoch+1))
# comment it out after debugging
if epoch>=24:
self.save_model(output_dir, epoch=(epoch+1), iteration=0, losses=losses)
## load the best weights
self.network.load_state_dict(best_model_wts)
# best epoch is 100
self.save_model(output_dir, epoch=100, iteration=0, losses=losses)
......@@ -2,6 +2,7 @@ from .CNNTrainer import CNNTrainer
from .MCCNNTrainer import MCCNNTrainer
from .DCGANTrainer import DCGANTrainer
from .ConditionalGANTrainer import ConditionalGANTrainer
from .FASNetTrainer import FASNetTrainer
# gets sphinx autodoc done right - don't remove it
__all__ = [_ for _ in dir() if not _.startswith('_')]
......
......@@ -71,6 +71,7 @@ setup(
'console_scripts' : [
'train_cnn.py = bob.learn.pytorch.scripts.train_cnn:main',
'train_mccnn.py = bob.learn.pytorch.scripts.train_mccnn:main',
'train_fasnet.py = bob.learn.pytorch.scripts.train_fasnet:main',
'train_dcgan.py = bob.learn.pytorch.scripts.train_dcgan:main',
'train_conditionalgan.py = bob.learn.pytorch.scripts.train_conditionalgan:main',
'train_network.py = bob.learn.pytorch.scripts.train_network:main',
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment