Commit e11e6883 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

[architectures, trainer] added LightCNN29 and LightCNN29v2 as well

parent 43fadd2a
Pipeline #26452 passed with stage
in 8 minutes and 12 seconds
......@@ -7,12 +7,14 @@ import torch.nn.functional as F
from .utils import MaxFeatureMap
from .utils import group
from .utils import resblock
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
"A light CNN for deep face representation with noisy labels", Wu, Xiang and He, Ran and Sun, Zhenan and Tan, Tieniu,
IEEE Transactions on Information Forensics and Security, vol 13, issue 11, 2018
Attributes
----------
......@@ -74,3 +76,175 @@ class LightCNN9(nn.Module):
out = self.fc2(x)
return out, x
class LightCNN29(nn.Module):
""" The class defining the light CNN with 29 layers
This class implements the CNN described in:
"A light CNN for deep face representation with noisy labels", Wu, Xiang and He, Ran and Sun, Zhenan and Tan, Tieniu,
IEEE Transactions on Information Forensics and Security, vol 13, issue 11, 2018
Attributes
----------
"""
def __init__(self, block=resblock, layers=[1, 2, 3, 4], num_classes=79077):
""" Init function
Parameters
----------
num_classes: int
The number of classes.
"""
super(LightCNN29, self).__init__()
self.conv1 = MaxFeatureMap(1, 48, 5, 1, 2)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
self.block1 = self._make_layer(block, layers[0], 48, 48)
self.group1 = group(48, 96, 3, 1, 1)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
self.block2 = self._make_layer(block, layers[1], 96, 96)
self.group2 = group(96, 192, 3, 1, 1)
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
self.block3 = self._make_layer(block, layers[2], 192, 192)
self.group3 = group(192, 128, 3, 1, 1)
self.block4 = self._make_layer(block, layers[3], 128, 128)
self.group4 = group(128, 128, 3, 1, 1)
self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
self.fc = MaxFeatureMap(8*8*128, 256, type=0)
self.fc2 = nn.Linear(256, num_classes)
def _make_layer(self, block, num_blocks, in_channels, out_channels):
"""
Parameters
----------
"""
layers = []
for i in range(0, num_blocks):
layers.append(block(in_channels, out_channels))
return nn.Sequential(*layers)
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.conv1(x)
x = self.pool1(x)
x = self.block1(x)
x = self.group1(x)
x = self.pool2(x)
x = self.block2(x)
x = self.group2(x)
x = self.pool3(x)
x = self.block3(x)
x = self.group3(x)
x = self.block4(x)
x = self.group4(x)
x = self.pool4(x)
x = x.view(x.size(0), -1)
fc = self.fc(x)
fc = F.dropout(fc, training=self.training)
out = self.fc2(fc)
return out, fc
class LightCNN29v2(nn.Module):
""" The class defining the light CNN with 29 layers (version 2)
This class implements the CNN described in:
"A light CNN for deep face representation with noisy labels", Wu, Xiang and He, Ran and Sun, Zhenan and Tan, Tieniu,
IEEE Transactions on Information Forensics and Security, vol 13, issue 11, 2018
Attributes
----------
"""
def __init__(self, block=resblock, layers=[1, 2, 3, 4], num_classes=79077):
""" Init function
Parameters
----------
num_classes: int
The number of classes.
"""
super(LightCNN29v2, self).__init__()
self.conv1 = MaxFeatureMap(1, 48, 5, 1, 2)
self.block1 = self._make_layer(block, layers[0], 48, 48)
self.group1 = group(48, 96, 3, 1, 1)
self.block2 = self._make_layer(block, layers[1], 96, 96)
self.group2 = group(96, 192, 3, 1, 1)
self.block3 = self._make_layer(block, layers[2], 192, 192)
self.group3 = group(192, 128, 3, 1, 1)
self.block4 = self._make_layer(block, layers[3], 128, 128)
self.group4 = group(128, 128, 3, 1, 1)
self.fc = nn.Linear(8*8*128, 256)
self.fc2 = nn.Linear(256, num_classes, bias=False)
def _make_layer(self, block, num_blocks, in_channels, out_channels):
"""
Parameters
----------
"""
layers = []
for i in range(0, num_blocks):
layers.append(block(in_channels, out_channels))
return nn.Sequential(*layers)
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.conv1(x)
x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2)
x = self.block1(x)
x = self.group1(x)
x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2)
x = self.block2(x)
x = self.group2(x)
x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2)
x = self.block3(x)
x = self.group3(x)
x = self.block4(x)
x = self.group4(x)
x = F.max_pool2d(x, 2) + F.avg_pool2d(x, 2)
x = x.view(x.size(0), -1)
fc = self.fc(x)
x = F.dropout(fc, training=self.training)
out = self.fc2(x)
return out, fc
from .CNN8 import CNN8
from .CASIANet import CASIANet
from .LightCNN import LightCNN9
from .LightCNN import LightCNN29
from .LightCNN import LightCNN29v2
from .DCGAN import DCGAN_generator
from .DCGAN import DCGAN_discriminator
......
......@@ -163,3 +163,19 @@ class group(nn.Module):
x = self.conv(x)
return x
class resblock(nn.Module):
""" Class implementing ...
"""
def __init__(self, in_channels, out_channels):
super(resblock, self).__init__()
self.conv1 = MaxFeatureMap(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
self.conv2 = MaxFeatureMap(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
def forward(self, x):
res = x
out = self.conv1(x)
out = self.conv2(out)
out = out + res
return out
......@@ -70,9 +70,11 @@ class CNNTrainer(object):
"""
try:
cp = torch.load(model_filename)
logger.info("model {} loaded".format(model_filename))
except RuntimeError:
# pre-trained model was probably saved using nn.DataParallel ...
cp = torch.load(model_filename, map_location='cpu')
logger.info("model {} loaded on CPU".format(model_filename))
if 'state_dict' in cp:
from collections import OrderedDict
......@@ -82,11 +84,14 @@ class CNNTrainer(object):
new_state_dict[name] = v
cp['state_dict'] = new_state_dict
logger.info("state_dict modified")
###########################################################################################################
### 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):
# LightCNN
if isinstance(self.network, bob.learn.pytorch.architectures.LightCNN.LightCNN9) \
or isinstance(self.network, bob.learn.pytorch.architectures.LightCNN.LightCNN29) \
or isinstance(self.network, bob.learn.pytorch.architectures.LightCNN.LightCNN29v2):
last_layer_weight = 'fc2.weight'
last_layer_bias = 'fc2.bias'
......@@ -99,9 +104,10 @@ class CNNTrainer(object):
var = 1.0 / (cp['state_dict'][last_layer_weight].shape[0])
np_weights = numpy.random.normal(loc=0.0, scale=var, size=((self.num_classes+1), cp['state_dict'][last_layer_weight].shape[1]))
cp['state_dict'][last_layer_weight] = torch.from_numpy(np_weights)
cp['state_dict'][last_layer_bias] = torch.zeros(((self.num_classes+1),))
#self.network.load_state_dict(cp['state_dict'], strict=False)
if not (isinstance(self.network, bob.learn.pytorch.architectures.LightCNN.LightCNN29v2)):
cp['state_dict'][last_layer_bias] = torch.zeros(((self.num_classes+1),))
self.network.load_state_dict(cp['state_dict'], strict=True)
logger.info("state_dict loaded for {} with {} classes".format(type(self.network), self.num_classes))
# CNN8
if isinstance(self.network, bob.learn.pytorch.architectures.CNN8):
......@@ -116,6 +122,7 @@ class CNNTrainer(object):
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)
logger.info("state_dict loaded for {} with {} classes".format(type(self.network), self.num_classes))
# CASIANet
if isinstance(self.network, bob.learn.pytorch.architectures.CASIANet):
......@@ -130,6 +137,7 @@ class CNNTrainer(object):
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)
logger.info("state_dict loaded for {} with {} classes".format(type(self.network), self.num_classes))
###########################################################################################################
......
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