Masterdqqq/Facial_Expression_Recognition
1
1"""2File: model.py3Author: Elena Ryumina and Dmitry Ryumin4Description: This module provides model architectures.5License: MIT License6"""7 8import torch9import torch.nn as nn10import torch.nn.functional as F11import math12 13class Bottleneck(nn.Module):14 expansion = 415 def __init__(self, in_channels, out_channels, i_downsample=None, stride=1):16 super(Bottleneck, self).__init__()17 18 self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, padding=0, bias=False)19 self.batch_norm1 = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.99)20 21 self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding='same', bias=False)22 self.batch_norm2 = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.99)23 24 self.conv3 = nn.Conv2d(out_channels, out_channels*self.expansion, kernel_size=1, stride=1, padding=0, bias=False)25 self.batch_norm3 = nn.BatchNorm2d(out_channels*self.expansion, eps=0.001, momentum=0.99)26 27 self.i_downsample = i_downsample28 self.stride = stride29 self.relu = nn.ReLU()30 31 def forward(self, x):32 identity = x.clone()33 x = self.relu(self.batch_norm1(self.conv1(x)))34 35 x = self.relu(self.batch_norm2(self.conv2(x)))36 37 x = self.conv3(x)38 x = self.batch_norm3(x)39 40 #downsample if needed41 if self.i_downsample is not None:42 identity = self.i_downsample(identity)43 #add identity44 x+=identity45 x=self.relu(x)46 47 return x48 49class Conv2dSame(torch.nn.Conv2d):50 51 def calc_same_pad(self, i: int, k: int, s: int, d: int) -> int:52 return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)53 54 def forward(self, x: torch.Tensor) -> torch.Tensor:55 ih, iw = x.size()[-2:]56 57 pad_h = self.calc_same_pad(i=ih, k=self.kernel_size[0], s=self.stride[0], d=self.dilation[0])58 pad_w = self.calc_same_pad(i=iw, k=self.kernel_size[1], s=self.stride[1], d=self.dilation[1])59 60 if pad_h > 0 or pad_w > 0:61 x = F.pad(62 x, [pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2]63 )64 return F.conv2d(65 x,66 self.weight,67 self.bias,68 self.stride,69 self.padding,70 self.dilation,71 self.groups,72 )73 74class ResNet(nn.Module):75 def __init__(self, ResBlock, layer_list, num_classes, num_channels=3):76 super(ResNet, self).__init__()77 self.in_channels = 6478 79 self.conv_layer_s2_same = Conv2dSame(num_channels, 64, 7, stride=2, groups=1, bias=False)80 self.batch_norm1 = nn.BatchNorm2d(64, eps=0.001, momentum=0.99)81 self.relu = nn.ReLU()82 self.max_pool = nn.MaxPool2d(kernel_size = 3, stride=2)83 84 self.layer1 = self._make_layer(ResBlock, layer_list[0], planes=64, stride=1)85 self.layer2 = self._make_layer(ResBlock, layer_list[1], planes=128, stride=2)86 self.layer3 = self._make_layer(ResBlock, layer_list[2], planes=256, stride=2)87 self.layer4 = self._make_layer(ResBlock, layer_list[3], planes=512, stride=2)88 89 self.avgpool = nn.AdaptiveAvgPool2d((1,1))90 self.fc1 = nn.Linear(512*ResBlock.expansion, 512)91 self.relu1 = nn.ReLU()92 self.fc2 = nn.Linear(512, num_classes)93 94 def extract_features(self, x):95 x = self.relu(self.batch_norm1(self.conv_layer_s2_same(x)))96 x = self.max_pool(x)97 # print(x.shape)98 x = self.layer1(x)99 x = self.layer2(x)100 x = self.layer3(x)101 x = self.layer4(x)102 103 x = self.avgpool(x)104 x = x.reshape(x.shape[0], -1)105 x = self.fc1(x)106 return x107 108 def forward(self, x):109 x = self.extract_features(x)110 x = self.relu1(x)111 x = self.fc2(x)112 return x113 114 def _make_layer(self, ResBlock, blocks, planes, stride=1):115 ii_downsample = None116 layers = []117 118 if stride != 1 or self.in_channels != planes*ResBlock.expansion:119 ii_downsample = nn.Sequential(120 nn.Conv2d(self.in_channels, planes*ResBlock.expansion, kernel_size=1, stride=stride, bias=False, padding=0),121 nn.BatchNorm2d(planes*ResBlock.expansion, eps=0.001, momentum=0.99)122 )123 124 layers.append(ResBlock(self.in_channels, planes, i_downsample=ii_downsample, stride=stride))125 self.in_channels = planes*ResBlock.expansion126 127 for i in range(blocks-1):128 layers.append(ResBlock(self.in_channels, planes))129 130 return nn.Sequential(*layers)131 132def ResNet50(num_classes, channels=3):133 return ResNet(Bottleneck, [3,4,6,3], num_classes, channels)134 135 136class LSTMPyTorch(nn.Module):137 def __init__(self):138 super(LSTMPyTorch, self).__init__()139 140 self.lstm1 = nn.LSTM(input_size=512, hidden_size=512, batch_first=True, bidirectional=False)141 self.lstm2 = nn.LSTM(input_size=512, hidden_size=256, batch_first=True, bidirectional=False)142 self.fc = nn.Linear(256, 7)143 self.softmax = nn.Softmax(dim=1)144 145 def forward(self, x):146 x, _ = self.lstm1(x)147 x, _ = self.lstm2(x) 148 x = self.fc(x[:, -1, :])149 x = self.softmax(x)150 return x