Commit 0f9a2e11 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

Merge branch 'lightCNN' into 'master'

Light cnn

See merge request !9
parents 1494e623 79b7432b
Pipeline #26412 failed with stages
in 6 minutes and 59 seconds
#!/usr/bin/env python
# encoding: utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
......
#!/usr/bin/env python
# encoding: utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from .utils import MaxFeatureMap
from .utils import group
class LightCNN9(nn.Module):
""" The class defining the light CNN with 9 layers
This class implements the CNN described in:
"Learning Face Representation From Scratch", D. Yi, Z. Lei, S. Liao and S.z. Li, 2014
Attributes
----------
features: :py:class:`torch.nn.Module`
The output of the convolutional / max layers
avgpool: :py:class:`torch.nn.Module`
The output of the average pooling layer (used as embedding)
classifier: :py:class:`torch.nn.Module`
The output of the last linear (logits)
"""
def __init__(self, num_classes=79077):
""" Init function
Parameters
----------
num_classes: int
The number of classes.
"""
super(LightCNN9, self).__init__()
self.features = nn.Sequential(
MaxFeatureMap(1, 48, 5, 1, 2),
nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True),
group(48, 96, 3, 1, 1),
nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True),
group(96, 192, 3, 1, 1),
nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True),
group(192, 128, 3, 1, 1),
group(128, 128, 3, 1, 1),
nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True),
)
self.fc1 = MaxFeatureMap(8*8*128, 256, type=0)
self.fc2 = nn.Linear(256, num_classes)
def forward(self, x):
""" Propagate data through the network
Parameters
----------
x: :py:class:`torch.Tensor`
The data to forward through the network. Image of size 1x128x128
Returns
-------
out: :py:class:`torch.Tensor`
class probabilities
x: :py:class:`torch.Tensor`
Output of the penultimate layer (i.e. embedding)
"""
x = self.features(x)
x = x.view(x.size(0), -1)
x = self.fc1(x)
x = F.dropout(x, training=self.training)
out = self.fc2(x)
return out, x
from .CNN8 import CNN8
from .CASIANet import CASIANet
from .LightCNN import LightCNN9
from .DCGAN import DCGAN_generator
from .DCGAN import DCGAN_discriminator
......
#!/usr/bin/env python
# encoding: utf-8
import torch
import torch.nn as nn
......@@ -51,3 +55,111 @@ def weights_init(m):
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
class MaxFeatureMap(nn.Module):
""" Class defining the max feature map
Attributes
----------
out_channels: int
the number of output channels ?
filter: either :py:class:`torch.nn.Conv2D` or :py:class:`torch.nn.Linear`
"""
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, type=1):
""" Init function
Parameters
----------
in_channels: int
the number of input channels
out_channels: int
the number of output channels
kernel_size: int
The size of the kernel in the convolution
stride: int
The stride in the convolution
padding: int
The padding (default to)
type: int
??
"""
super(MaxFeatureMap, self).__init__()
self.out_channels = out_channels
if type == 1:
self.filter = nn.Conv2d(in_channels, 2*out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
else:
self.filter = nn.Linear(in_channels, 2*out_channels)
def forward(self, x):
""" Forward function
Propagates data through the Max Feature Map
Parameters
----------
x: :py:class:`torch.Tensor`
The data to forward through the MFM
Returns
-------
py:class:`torch.Tensor`
"""
x = self.filter(x)
out = torch.split(x, self.out_channels, 1)
return torch.max(out[0], out[1])
class group(nn.Module):
""" Class implementing ...
Attributes
----------
"""
def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
""" Init function
Parameters
----------
in_channels: int
the number of input channels
out_channels: int
the number of output channels
kernel_size: int
The size of the kernel in the convolution
stride: int
The stride in the convolution
padding: int
The padding (default to)
"""
super(group, self).__init__()
self.conv_a = MaxFeatureMap(in_channels, in_channels, 1, 1, 0)
self.conv = MaxFeatureMap(in_channels, out_channels, kernel_size, stride, padding)
def forward(self, x):
""" Forward function
Propagates data through the Max Feature Map
Parameters
----------
x: :py:class:`torch.Tensor`
The data to forward through the MFM
Returns
-------
py:class:`torch.Tensor`
"""
x = self.conv_a(x)
x = self.conv(x)
return x
......@@ -8,6 +8,7 @@ from .utils import RollChannels
from .utils import ToTensor
from .utils import Normalize
from .utils import Resize
from .utils import ToGray
from .utils import map_labels
from .utils import ConcatDataset
......
......@@ -55,9 +55,10 @@ class CasiaWebFaceDataset(Dataset):
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):
"""Returns the length of the dataset (i.e. nb of examples)
......
......@@ -67,6 +67,19 @@ class Resize(object):
sample['image'] = sample['image'][..., numpy.newaxis]
return sample
class ToGray(object):
def __init__(self):
self.op = transforms.Grayscale()
def __call__(self, sample):
# convert to PIL image
from PIL.Image import fromarray
img = fromarray(sample['image'].squeeze())
img = self.op(img)
sample['image'] = numpy.array(img)
sample['image'] = sample['image'][..., numpy.newaxis]
return sample
def map_labels(raw_labels, start_index=0):
"""
......@@ -77,18 +90,9 @@ def map_labels(raw_labels, start_index=0):
for i in range(len(possible_labels)):
l = possible_labels[i]
labels[numpy.where(labels==l)[0][0]] = i + start_index
# -----
# map back to native int, resolve the problem with dataset concatenation
# it does: line 78 is now ok
# for some reason, it was not working when the type of id labels were numpy.int64 ...
labels_int = []
for i in range(len(labels)):
labels_int.append(labels[i].item())
return labels_int
labels[numpy.where(labels==l)[0]] = i + start_index
return labels
from torch.utils.data import Dataset
import bob.io.base
......
......@@ -103,14 +103,15 @@ def main(user_input=None):
# get data
if hasattr(configuration, 'dataset'):
dataloader = torch.utils.data.DataLoader(configuration.dataset, batch_size=batch_size, shuffle=True)
logger.info("There are {} training images from {} categories".format(len(configuration.dataset), numpy.max(configuration.dataset.id_labels)))
num_classes = numpy.max(configuration.dataset.id_labels)
logger.info("There are {} training images from {} categories".format(len(configuration.dataset), num_classes))
else:
logger.error("Please provide a dataset in your configuration file !")
sys.exit()
# train the network
if hasattr(configuration, 'network'):
trainer = CNNTrainer(configuration.network, batch_size=batch_size, use_gpu=use_gpu, verbosity_level=verbosity_level)
trainer = CNNTrainer(configuration.network, batch_size=batch_size, use_gpu=use_gpu, verbosity_level=verbosity_level, num_classes=num_classes)
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 !")
......
......@@ -33,6 +33,15 @@ def test_architectures():
assert output.shape == torch.Size([1, 20])
assert emdedding.shape == torch.Size([1, 512])
# LightCNN9
a = numpy.random.rand(1, 1, 128, 128).astype("float32")
t = torch.from_numpy(a)
from ..architectures import LightCNN9
net = LightCNN9()
output, emdedding = net.forward(t)
assert output.shape == torch.Size([1, 79077])
assert emdedding.shape == torch.Size([1, 256])
# DCGAN
d = numpy.random.rand(1, 3, 64, 64).astype("float32")
t = torch.from_numpy(d)
......@@ -100,12 +109,18 @@ def test_map_labels():
labels = ['1', '4', '7']
from ..datasets import map_labels
new_labels = map_labels(labels)
new_labels = sorted(new_labels)
assert new_labels == ['0', '1', '2']
assert '0' in new_labels, "new_labels = {}".format(new_labels)
assert '1' in new_labels, "new_labels = {}".format(new_labels)
assert '2' in new_labels, "new_labels = {}".format(new_labels)
#new_labels = sorted(new_labels)
#assert new_labels == ['0', '1', '2']
new_labels = map_labels(labels, start_index = 5)
new_labels = sorted(new_labels)
assert new_labels == ['5', '6', '7']
#new_labels = sorted(new_labels)
assert '5' in new_labels, "new_labels = {}".format(new_labels)
assert '6' in new_labels, "new_labels = {}".format(new_labels)
assert '7' in new_labels, "new_labels = {}".format(new_labels)
#assert new_labels == ['5', '6', '7']
from torch.utils.data import Dataset
......@@ -115,7 +130,7 @@ class DummyDataSet(Dataset):
def __len__(self):
return 100
def __getitem__(self, idx):
data = numpy.random.rand(3, 128, 128).astype("float32")
data = numpy.random.rand(1, 128, 128).astype("float32")
label = numpy.random.randint(20)
sample = {'image': torch.from_numpy(data), 'label': label}
return sample
......@@ -123,8 +138,8 @@ class DummyDataSet(Dataset):
def test_CNNtrainer():
from ..architectures import CNN8
net = CNN8(20)
from ..architectures import LightCNN9
net = LightCNN9(20)
dataloader = torch.utils.data.DataLoader(DummyDataSet(), batch_size=32, shuffle=True)
......
......@@ -13,6 +13,7 @@ logger = bob.core.log.setup("bob.learn.pytorch")
import time
import os
import numpy
class CNNTrainer(object):
"""
......@@ -31,7 +32,7 @@ class CNNTrainer(object):
"""
def __init__(self, network, batch_size=64, use_gpu=False, verbosity_level=2):
def __init__(self, network, batch_size=64, use_gpu=False, verbosity_level=2, num_classes=2):
""" Init function
Parameters
......@@ -44,9 +45,12 @@ class CNNTrainer(object):
If you would like to use the gpu
verbosity_level: int
The level of verbosity output to stdout
num_classes: int
The number of classes
"""
self.network = network
self.num_classes = num_classes
self.batch_size = batch_size
self.use_gpu = use_gpu
self.criterion = nn.CrossEntropyLoss()
......@@ -56,31 +60,88 @@ class CNNTrainer(object):
bob.core.log.set_verbosity_level(logger, verbosity_level)
def load_model(self, model_filename):
"""Loads an existing model
def load_and_initialize_model(self, model_filename):
""" Loads and initialize a 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
model_filename: str
"""
cp = torch.load(model_filename)
self.network.load_state_dict(cp['state_dict'])
start_epoch = cp['epoch']
start_iter = cp['iteration']
losses = cp['loss']
try:
cp = torch.load(model_filename)
#self.network.load_state_dict(cp['state_dict'])
except RuntimeError:
# pre-trained model was probably saved using nn.DataParallel ...
cp = torch.load(model_filename, map_location='cpu')
if 'state_dict' in cp:
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in cp['state_dict'].items():
name = k[7:]
new_state_dict[name] = v
cp['state_dict'] = new_state_dict
print(type(self.network))
###########################################################################################################
### for each defined architecture, get the output size in pre-trained model, and change it if necessary ###
# LightCNN9
if isinstance(self.network, bob.learn.pytorch.architectures.LightCNN.LightCNN9):
num_classes_pretrained = cp['state_dict']['fc2.weight'].shape[0]
if num_classes_pretrained == self.num_classes:
self.network.load_state_dict(cp['state_dict'])
else:
var = 1.0 / (cp['state_dict']['fc2.weight'].shape[0])
np_weights = numpy.random.normal(loc=0.0, scale=var, size=((self.num_classes+1), cp['state_dict']['fc2.weight'].shape[1]))
cp['state_dict']['fc2.weight'] = torch.from_numpy(np_weights)
cp['state_dict']['fc2.bias'] = torch.zeros(((self.num_classes+1),))
#self.network.load_state_dict(cp['state_dict'], strict=False)
self.network.load_state_dict(cp['state_dict'], strict=True)
# CNN8
if isinstance(self.network, bob.learn.pytorch.architectures.CNN8):
num_classes_pretrained = cp['state_dict']['classifier.weight'].shape[0]
if num_classes_pretrained == self.num_classes:
self.network.load_state_dict(cp['state_dict'])
else:
var = 1.0 / (cp['state_dict']['classifier.weight'].shape[0])
np_weights = numpy.random.normal(loc=0.0, scale=var, size=((self.num_classes+1), cp['state_dict']['classifier.weight'].shape[1]))
cp['state_dict']['classifier.weight'] = torch.from_numpy(np_weights)
cp['state_dict']['classifier.bias'] = torch.zeros(((self.num_classes+1),))
#self.network.load_state_dict(cp['state_dict'], strict=False)
self.network.load_state_dict(cp['state_dict'], strict=True)
# CASIANet
if isinstance(self.network, bob.learn.pytorch.architectures.CASIANet):
num_classes_pretrained = cp['state_dict']['classifier.weight'].shape[0]
if num_classes_pretrained == self.num_classes:
self.network.load_state_dict(cp['state_dict'])
else:
var = 1.0 / (cp['state_dict']['classifier.weight'].shape[0])
np_weights = numpy.random.normal(loc=0.0, scale=var, size=((self.num_classes+1), cp['state_dict']['classifier.weight'].shape[1]))
cp['state_dict']['classifier.weight'] = torch.from_numpy(np_weights)
cp['state_dict']['classifier.bias'] = torch.zeros(((self.num_classes+1),))
#self.network.load_state_dict(cp['state_dict'], strict=False)
self.network.load_state_dict(cp['state_dict'], strict=True)
###########################################################################################################
start_epoch = 0
start_iter = 0
losses = []
if 'epoch' in cp.keys():
start_epoch = cp['epoch']
if 'iteration' in cp.keys():
start_iter = cp['iteration']
if 'losses' in cp.keys():
losses = cp['epoch']
return start_epoch, start_iter, losses
......@@ -133,8 +194,15 @@ class CNNTrainer(object):
# 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]))
start_epoch, start_iter, losses = self.load_and_initialize_model(model)
if start_epoch != 0:
logger.info('Previous network was trained up to epoch {}, iteration {}'.format(start_epoch, start_iter, losses[-1]))
if losses:
logger.info('Last loss = {}'.format(losses[-1]))
else:
logger.info('Starting training / fine-tuning from pre-trained model')
else:
start_epoch = 0
start_iter = 0
......
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