Commit 1f61fbf1 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH

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

parent 6515dc97
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
from .DCGAN import weights_init
from .CNN8 import CNN8
# gets sphinx autodoc done right - don't remove it
__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 .casia_webface import CasiaDataset
from .casia_webface import CasiaWebFaceDataset
from .fargo import FargoDataset
# transforms
from .utils import FaceCropper
from .utils import RollChannels
from .utils import ToTensor
from .utils import Normalize
......
......@@ -8,14 +8,14 @@ import numpy
from torch.utils.data import Dataset, DataLoader
import bob.db.casia_webface
import bob.io.base
import bob.io.image
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
......@@ -30,9 +30,58 @@ class CasiaDataset(Dataset):
transform: torchvision.transforms
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):
self.root_dir = root_dir
self.transform = transform
......@@ -75,7 +124,7 @@ class CasiaDataset(Dataset):
return the length of the dataset (i.e. nb of examples)
"""
return len(self.data_files)
def __getitem__(self, idx):
"""
......
......@@ -5,6 +5,28 @@ import numpy
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):
"""
......@@ -95,3 +117,5 @@ class ConcatDataset(Dataset):
sample = self.transform(sample)
return sample
This diff is collapsed.
#!/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
from .ConditionalGANTrainer import ConditionalGANTrainer
from .ImprovedWassersteinCGANTrainer import IWCGANTrainer
from .DRGANTrainer import DRGANTrainer
from .CNNTrainer import CNNTrainer
# gets sphinx autodoc done right - don't remove it
__all__ = [_ for _ in dir() if not _.startswith('_')]
......
......@@ -82,6 +82,8 @@ setup(
'show_training_images.py = bob.learn.pytorch.scripts.show_training_images:main',
'show_training_stats.py = bob.learn.pytorch.scripts.show_training_stats: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',
],
......
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