Skip to content
Snippets Groups Projects
Commit 1f61fbf1 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

[all] added stuff to train Xiaojiang CNN8 model (architecture, trainer, new Casia dataset)

parent 6515dc97
Branches
Tags
No related merge requests found
import torch
import torch.nn as nn
import torch.nn.functional as F
from .utils import make_conv_layers
CNN8_CONFIG = [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M']
class CNN8(nn.Module):
def __init__(self, num_cls, drop_rate=0.5):
super(CNN8, self).__init__()
self.num_classes = num_cls
self.drop_rate = float(drop_rate)
self.conv = make_conv_layers(CNN8_CONFIG)
self.avgpool = nn.AvgPool2d(8)
self.classifier = nn.Linear(512, self.num_classes)
def forward(self, x):
x = self.conv(x)
x = self.avgpool(x)
x = x.view(x.size(0), -1)
x = F.dropout(x, p = self.drop_rate, training=self.training)
out = self.classifier(x)
return out, x # x for feature
...@@ -17,6 +17,8 @@ from .DRGANOriginal import DRGANOriginal_discriminator ...@@ -17,6 +17,8 @@ from .DRGANOriginal import DRGANOriginal_discriminator
from .DCGAN import weights_init from .DCGAN import weights_init
from .CNN8 import CNN8
# gets sphinx autodoc done right - don't remove it # gets sphinx autodoc done right - don't remove it
__all__ = [_ for _ in dir() if not _.startswith('_')] __all__ = [_ for _ in dir() if not _.startswith('_')]
import torch
import torch.nn as nn
def make_conv_layers(cfg, input_c = 3):
layers = []
in_channels = input_c
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
layers += [conv2d, nn.ReLU()]
in_channels = v
return nn.Sequential(*layers)
from .multipie import MultiPIEDataset from .multipie import MultiPIEDataset
from .casia_webface import CasiaDataset from .casia_webface import CasiaDataset
from .casia_webface import CasiaWebFaceDataset
from .fargo import FargoDataset
# transforms # transforms
from .utils import FaceCropper
from .utils import RollChannels from .utils import RollChannels
from .utils import ToTensor from .utils import ToTensor
from .utils import Normalize from .utils import Normalize
......
...@@ -8,14 +8,14 @@ import numpy ...@@ -8,14 +8,14 @@ import numpy
from torch.utils.data import Dataset, DataLoader from torch.utils.data import Dataset, DataLoader
import bob.db.casia_webface
import bob.io.base import bob.io.base
import bob.io.image import bob.io.image
from .utils import map_labels from .utils import map_labels
class CasiaDataset(Dataset):
"""Casia WebFace dataset. class CasiaWebFaceDataset(Dataset):
"""Casia WebFace dataset (for CNN training).
Class representing the CASIA WebFace dataset Class representing the CASIA WebFace dataset
...@@ -30,9 +30,58 @@ class CasiaDataset(Dataset): ...@@ -30,9 +30,58 @@ class CasiaDataset(Dataset):
transform: torchvision.transforms transform: torchvision.transforms
The transform(s) to apply to the face images The transform(s) to apply to the face images
""" """
def __init__(self, root_dir, transform=None, start_index=0):
self.root_dir = root_dir
self.transform = transform
self.data_files = []
id_labels = []
for root, dirs, files in os.walk(self.root_dir):
for name in files:
filename = os.path.split(os.path.join(root, name))[-1]
path = root.split(os.sep)
subject = int(path[-1])
self.data_files.append(os.path.join(root, name))
id_labels.append(subject)
self.id_labels = map_labels(id_labels, start_index)
def __len__(self):
"""
return the length of the dataset (i.e. nb of examples)
"""
return len(self.data_files)
def __getitem__(self, idx):
"""
return a sample from the dataset
"""
image = bob.io.base.load(self.data_files[idx])
identity = self.id_labels[idx]
sample = {'image': image, 'id': identity}
if self.transform:
sample = self.transform(sample)
return sample
class CasiaDataset(Dataset):
"""Casia WebFace dataset.
Class representing the CASIA WebFace dataset
# TODO: Start from original data and annotations - Guillaume HEUSCH, 06-11-2017 **Parameters**
root-dir: path
The path to the data
frontal_only: boolean
If you want to only use frontal faces
transform: torchvision.transforms
The transform(s) to apply to the face images
"""
def __init__(self, root_dir, frontal_only=False, transform=None, start_index=0): def __init__(self, root_dir, frontal_only=False, transform=None, start_index=0):
self.root_dir = root_dir self.root_dir = root_dir
self.transform = transform self.transform = transform
...@@ -75,7 +124,7 @@ class CasiaDataset(Dataset): ...@@ -75,7 +124,7 @@ class CasiaDataset(Dataset):
return the length of the dataset (i.e. nb of examples) return the length of the dataset (i.e. nb of examples)
""" """
return len(self.data_files) return len(self.data_files)
def __getitem__(self, idx): def __getitem__(self, idx):
""" """
......
...@@ -5,6 +5,28 @@ import numpy ...@@ -5,6 +5,28 @@ import numpy
import torchvision.transforms as transforms import torchvision.transforms as transforms
class FaceCropper():
"""
Class to crop a face, based on eyes position
"""
def __init__(self, cropped_height, cropped_width):
# the face cropper
from bob.bio.face.preprocessor import FaceCrop
cropped_image_size = (cropped_height, cropped_width)
right_eye_pos = (cropped_height // 5, cropped_width // 4 -1)
left_eye_pos = (cropped_height // 5, cropped_width // 4 * 3)
cropped_positions = {'leye': left_eye_pos, 'reye': right_eye_pos}
self.face_cropper = FaceCrop(cropped_image_size=cropped_image_size,
cropped_positions=cropped_positions,
color_channel='rgb',
dtype='uint8'
)
def __call__(self, sample):
cropped = self.face_cropper(sample['image'], sample['eyes'])
sample['image'] = cropped
return sample
class RollChannels(object): class RollChannels(object):
""" """
...@@ -95,3 +117,5 @@ class ConcatDataset(Dataset): ...@@ -95,3 +117,5 @@ class ConcatDataset(Dataset):
sample = self.transform(sample) sample = self.transform(sample)
return sample return sample
#!/usr/bin/env python
# encoding: utf-8
""" Train a CNN for face recognition
Usage:
%(prog)s [--net=<string>] [--model=<string>] [--dropout=<float>] [--learning-rate=<float>]
[--batch-size=<int>] [--epochs=<int>] [--nclasses=<int>]
[--output-dir=<path>] [--use-gpu] [--seed=<int>] [--verbose ...]
Options:
-h, --help Shows this help message and exits
--net=<string> Name of a network [default: CNN8]
--model=<string> Filename of the model to load.
--dropout=<float> Dropout rate [default: 0.5]
--learning-rate=<float> Learning rate [default: 0.01]
--batch-size=<int> Batch size [default: 64]
--epochs=<int> Number of training epochs [default: 20]
--nclasses=<int> Number of classes [default: 10575]
-o, --output-dir=<path> Dir to save stuff [default: CNN8-training]
-g, --use-gpu Use the GPU
-S, --seed=<int> The random seed [default: 3]
-v, --verbose Increase the verbosity (may appear multiple times).
Example:
To run the training process
$ %(prog)s --net CNN8 --batch-size 64 --epochs 25
See '%(prog)s --help' for more information.
"""
import os, sys
import pkg_resources
import bob.core
logger = bob.core.log.setup("bob.learn.pytorch")
from docopt import docopt
version = pkg_resources.require('bob.learn.pytorch')[0].version
import numpy
import bob.io.base
import time
from pdb import set_trace as st
import atexit
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision.transforms as transforms
from bob.learn.pytorch.datasets import CasiaWebFaceDataset
from bob.learn.pytorch.datasets import RollChannels
from bob.learn.pytorch.datasets import ToTensor
from bob.learn.pytorch.datasets import Normalize
from bob.learn.pytorch.trainers import CNNTrainer
#from lib import networks,folder
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 CNN (%s)' % version,)
# verbosity
verbosity_level = args['--verbose']
bob.core.log.set_verbosity_level(logger, verbosity_level)
# get the arguments
network = args['--net']
model = args['--model']
dropout = float(args['--dropout'])
learning_rate = float(args['--learning-rate'])
batch_size = int(args['--batch-size'])
epochs = int(args['--epochs'])
number_of_classes = int(args['--nclasses'])
output_dir = args['--output-dir']
use_gpu = bool(args['--use-gpu'])
seed = int(args['--seed'])
# 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")
bob.io.base.create_directories_safe(output_dir)
# ============
# === DATA ===
# ============
# WARNING with the transforms ... act on labels too, at some point, I may have to write my own
# Also, in 'ToTensor', there is a reshape performed from: HxWxC to CxHxW
face_dataset = CasiaWebFaceDataset(root_dir='/idiap/project/fargo/xpeng_prepro/CASIA-Webface-crop-128/',
transform=transforms.Compose([
RollChannels(), # bob to skimage:
ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
)
dataloader = torch.utils.data.DataLoader(face_dataset, batch_size=batch_size, shuffle=True)
logger.info("There are {} training images from {} identities".format(len(face_dataset), numpy.max(face_dataset.id_labels)))
# =======
# NETWORK
# =======
if network == 'CNN8':
from bob.learn.pytorch.architectures import CNN8
net = CNN8(number_of_classes, dropout)
print(net)
#print model
#if args['--model']:
# cp = torch.load(args['--model'])
# net.load_state_dict(cp['state_dict'])
# start_epoch = cp['epoch']
# epoch_ = start_epoch
# start_iter = cp['iteration']
# losses = cp['loss']
# =====
# TRAIN
# =====
trainer = CNNTrainer(net, batch_size=batch_size, use_gpu=use_gpu, verbosity_level=verbosity_level)
trainer.train(dataloader, n_epochs=epochs, learning_rate=learning_rate, output_dir=output_dir, model=model)
#elif modname=='CASIA_NET':
# net = networks.CASIA_NET(int(args['--num_classes']), args['--drop'])
#elif modname=='ResNet26':
# net = networks.ResNet26(int(args['--num_classes']), args['--drop'])
#def exit_handler():
# if 'args' in globals():
# global losses, epoch_, iteration_, net
# print 'Ctrl-C! Safely exit after saving your model...'
# save_network(args['<save_dir>'], net, args['--net'], losses, epoch_, iteration_)
#
#def get_time_str():
# return time.strftime("%Y-%m-%d, %H:%M:%S ", time.localtime((time.time()) ))
#
#def print_info(msg):
# print get_time_str(), msg
# sys.stdout.flush()
#
#def setup(modname):
# global net
# if modname=='CNN8':
# net = networks.CNN8(int(args['--num_classes']), args['--drop'])
# elif modname=='CASIA_NET':
# net = networks.CASIA_NET(int(args['--num_classes']), args['--drop'])
# elif modname=='ResNet26':
# net = networks.ResNet26(int(args['--num_classes']), args['--drop'])
#
#
#def save_network(save_dir, network,network_label, trlosses = None, epoch=0, iteration = 0):
# save_filename = '{}_{}_{}.pth'.format(network_label,epoch, iteration)
# save_path = os.path.join(save_dir, save_filename)
# print_info('saving model to {}'.format(save_path))
# cp = {'epoch': epoch, 'iteration': iteration,
# 'loss': trlosses, 'state_dict': network.cpu().state_dict()}
# torch.save(cp, save_path)
# print_info('Successful!')
# if torch.cuda.is_available(): #back to gpu
# network.cuda()
#
#atexit.register(exit_handler)
#
#args = docopt.docopt(__doc__)
#if args['--visdom']:
# import visdom
# vis = visdom.Visdom()
#
#
#
#global losses, epoch_, iteration_, net
#if __name__ == '__main__':
# print args
# setup(args['--net']) # setup the net to train
# net.train()
#
# save_dir = args['<save_dir>']
# if not os.path.exists(save_dir):
# os.makedirs(save_dir)
#
# transform = transforms.Compose([
# transforms.RandomHorizontalFlip(),
# transforms.Scale((128,128)), # fix scale for all model
# transforms.ToTensor(),
# transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# dset = folder.ImageFolder(args['<data_dir>'], transform, filelist=args['--trainlist'])
# loader = data.DataLoader(dset,
# batch_size=int(args['--batch_size']), shuffle=True,
# num_workers=int(args['--num_workers']), pin_memory=args['--pm'])
# iters_per_epoch = len(loader)
# print_info('# training images = %d' % len(dset))
# print_info('# classes = %d' % len(dset.classes))
# print_info('iters_per_epoch = %d' % iters_per_epoch)
#
# if args['--gpu']: net.cuda()
# start_epoch = 0
# start_iter = 0
# iteration_ = 0
# epoch_ = 0
# losses = []
# if args['--resume']:
# cp = torch.load(args['--resume'])
# net.load_state_dict(cp['state_dict'])
# start_epoch = cp['epoch']
# epoch_ = start_epoch
# start_iter = cp['iteration']
# losses = cp['loss']
#
# criterion = nn.CrossEntropyLoss()
# optimizer = optim.SGD(net.parameters(), float(args['--lr']) , momentum = 0.9, weight_decay = 0.0005)
# # normally you watch visdom, make early stop, set smaller lr and resume
# # fisrt to finish the last end epoch if necessary
# if start_iter>0:
# if args['--visdom']:
# win = vis.line(X=np.arange(len(losses))*int(args['--plot_freq_iter']),Y=np.array(losses)) # see http://localhost:8097/
# for i, (input, target) in enumerate(loader):
# iteration_ = i
# if i>=start_iter:
# if args['--gpu']:
# input_var = torch.autograd.Variable(input.cuda())
# target_var = torch.autograd.Variable(target.cuda())
# else:
# input_var = torch.autograd.Variable(input)
# target_var = torch.autograd.Variable(target)
# output, _ = net(input_var)
# loss = criterion(output, target_var)
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
#
# if i % int(args['--plot_freq_iter']) == 0:
# print_info('Epoch: {}/{}, iter: {}/{}, loss: {}'.format(start_epoch, args['--epochs'], i, iters_per_epoch, loss.data[0]))
# losses.append(loss.data[0])
# if args['--visdom']:
# vis.line(X=np.arange(len(losses))*int(args['--plot_freq_iter']), Y=np.array(losses),
# opts={
# 'title':' loss over time',
# 'xlabel': 'iteration',
# 'ylabel': 'loss'},
# win=win)
# if i % int(args['--save_freq_iter']) == 0: #i!=0
# save_network(save_dir, net, args['--net'], losses, start_epoch, i)
# epoch_ += 1
# # now normal epochs
# for epoch in range(start_epoch+1, int(args['--epochs'])):
# epoch_ = epoch # for exit writing
# for i, (input, target) in enumerate(loader):
# iteration_ = i
# if args['--gpu']:
# input_var = torch.autograd.Variable(input.cuda())
# target_var = torch.autograd.Variable(target.cuda())
# else:
# input_var = torch.autograd.Variable(input)
# target_var = torch.autograd.Variable(target)
# output, _ = net(input_var)
# loss = criterion(output, target_var)
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
#
# if i % int(args['--plot_freq_iter']) == 0:
# print_info('Epoch: {}/{}, iter: {}/{}, loss: {}'.format(epoch, args['--epochs'], i, iters_per_epoch, loss.data[0]))
# losses.append(loss.data[0])
# if args['--visdom']:
# if i==0 and start_iter==0:
# win = vis.line(X=np.arange(len(losses))*int(args['--plot_freq_iter']),Y=np.array(losses)) # see http://localhost:8097/
# else:
# vis.line(X=np.arange(len(losses))*int(args['--plot_freq_iter']), Y=np.array(losses),
# opts={
# 'title':' loss over time',
# 'xlabel': 'iteration',
# 'ylabel': 'loss'},
# win=win)
# if i>0 and i % int(args['--save_freq_iter']) == 0: #i!=0
# save_network(save_dir, net, args['--net'], losses, epoch, i)
#
# if epoch>0 and epoch % int(args['--save_freq_epoch']) == 0:
# save_network(save_dir, net, args['--net'], losses, epoch, 0)
#
# print 'Done'
#!/usr/bin/env python
# encoding: utf-8
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
import torchvision.utils as vutils
import bob.core
logger = bob.core.log.setup("bob.learn.pytorch")
import time
class CNNTrainer(object):
"""
Class to train a CNN
**Parameters**
network: pytorch nn.Module
The network
batch_size: int
The size of your minibatch
use_gpu: boolean
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):
self.network = network
self.batch_size = batch_size
self.use_gpu = use_gpu
self.criterion = nn.CrossEntropyLoss()
if self.use_gpu:
self.network.cuda()
bob.core.log.set_verbosity_level(logger, verbosity_level)
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
The list of losses wfrom 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(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
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(save_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=20, learning_rate=0.01, output_dir='out', model=None):
"""
Function that performs the training.
**Parameters**
dataloader: pytorch 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 SGD optimizer
output_dir: path
The directory where you would like to save models
"""
# if model exists, load it
if model is not None:
start_epoch, start_iter, losses = self.load_model(model)
else:
start_epoch = 0
start_iter = 0
losses = []
# setup optimizer
optimizer = optim.SGD(self.network.parameters(), learning_rate, momentum = 0.9, weight_decay = 0.0005)
# let's go
for epoch in range((start_epoch, n_epochs):
for i, data in enumerate(dataloader, 0):
if i >= start_iter:
start = time.time()
images = data['image']
labels = data['id']
batch_size = len(images)
if self.use_gpu:
images = images.cuda()
labels = labels.cuda()
imagesv = Variable(images)
labelsv = Variable(labels)
output, _ = self.network(imagesv)
loss = self.criterion(output, labelsv)
optimizer.zero_grad()
loss.backward()
optimizer.step()
end = time.time()
logger.info("[{}/{}][{}/{}] => Loss = {} (time spent: {})".format(epoch, n_epochs, i, len(dataloader), loss.data[0], (end-start)))
losses.append(loss.data[0])
# do stuff - like saving models
logger.info("EPOCH {} DONE")
self.save_model(output_dir, epoch=(epoch+1), iteration=0, losses=losses)
...@@ -2,6 +2,7 @@ from .DCGANTrainer import DCGANTrainer ...@@ -2,6 +2,7 @@ from .DCGANTrainer import DCGANTrainer
from .ConditionalGANTrainer import ConditionalGANTrainer from .ConditionalGANTrainer import ConditionalGANTrainer
from .ImprovedWassersteinCGANTrainer import IWCGANTrainer from .ImprovedWassersteinCGANTrainer import IWCGANTrainer
from .DRGANTrainer import DRGANTrainer from .DRGANTrainer import DRGANTrainer
from .CNNTrainer import CNNTrainer
# gets sphinx autodoc done right - don't remove it # gets sphinx autodoc done right - don't remove it
__all__ = [_ for _ in dir() if not _.startswith('_')] __all__ = [_ for _ in dir() if not _.startswith('_')]
......
...@@ -82,6 +82,8 @@ setup( ...@@ -82,6 +82,8 @@ setup(
'show_training_images.py = bob.learn.pytorch.scripts.show_training_images:main', 'show_training_images.py = bob.learn.pytorch.scripts.show_training_images:main',
'show_training_stats.py = bob.learn.pytorch.scripts.show_training_stats:main', 'show_training_stats.py = bob.learn.pytorch.scripts.show_training_stats:main',
'sample_drgan.py = bob.learn.pytorch.scripts.sample_drgan:main', 'sample_drgan.py = bob.learn.pytorch.scripts.sample_drgan:main',
'test_fargo.py = bob.learn.pytorch.scripts.test_fargo:main',
'train_cnn.py = bob.learn.pytorch.scripts.train_cnn:main',
], ],
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment