Siamese network for image similarity estimation. The network is composed of two identical networks, one for each input. The output of each network is concatenated and passed to a linear layer. The output of the linear layer passed through a sigmoid function.
| 14 | |
| 15 | |
| 16 | class SiameseNetwork(nn.Module): |
| 17 | """ |
| 18 | Siamese network for image similarity estimation. |
| 19 | The network is composed of two identical networks, one for each input. |
| 20 | The output of each network is concatenated and passed to a linear layer. |
| 21 | The output of the linear layer passed through a sigmoid function. |
| 22 | `"FaceNet" <https://arxiv.org/pdf/1503.03832.pdf>`_ is a variant of the Siamese network. |
| 23 | This implementation varies from FaceNet as we use the `ResNet-18` model from |
| 24 | `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_ as our feature extractor. |
| 25 | In addition, we aren't using `TripletLoss` as the MNIST dataset is simple, so `BCELoss` can do the trick. |
| 26 | """ |
| 27 | def __init__(self): |
| 28 | super(SiameseNetwork, self).__init__() |
| 29 | # get resnet model |
| 30 | self.resnet = torchvision.models.resnet18(weights=None) |
| 31 | |
| 32 | # over-write the first conv layer to be able to read MNIST images |
| 33 | # as resnet18 reads (3,x,x) where 3 is RGB channels |
| 34 | # whereas MNIST has (1,x,x) where 1 is a gray-scale channel |
| 35 | self.resnet.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) |
| 36 | self.fc_in_features = self.resnet.fc.in_features |
| 37 | |
| 38 | # remove the last layer of resnet18 (linear layer which is before avgpool layer) |
| 39 | self.resnet = torch.nn.Sequential(*(list(self.resnet.children())[:-1])) |
| 40 | |
| 41 | # add linear layers to compare between the features of the two images |
| 42 | self.fc = nn.Sequential( |
| 43 | nn.Linear(self.fc_in_features * 2, 256), |
| 44 | nn.ReLU(inplace=True), |
| 45 | nn.Linear(256, 1), |
| 46 | ) |
| 47 | |
| 48 | self.sigmoid = nn.Sigmoid() |
| 49 | |
| 50 | # initialize the weights |
| 51 | self.resnet.apply(self.init_weights) |
| 52 | self.fc.apply(self.init_weights) |
| 53 | |
| 54 | def init_weights(self, m): |
| 55 | if isinstance(m, nn.Linear): |
| 56 | torch.nn.init.xavier_uniform_(m.weight) |
| 57 | m.bias.data.fill_(0.01) |
| 58 | |
| 59 | def forward_once(self, x): |
| 60 | output = self.resnet(x) |
| 61 | output = output.view(output.size()[0], -1) |
| 62 | return output |
| 63 | |
| 64 | def forward(self, input1, input2): |
| 65 | # get two images' features |
| 66 | output1 = self.forward_once(input1) |
| 67 | output2 = self.forward_once(input2) |
| 68 | |
| 69 | # concatenate both images' features |
| 70 | output = torch.cat((output1, output2), 1) |
| 71 | |
| 72 | # pass the concatenation to the linear layers |
| 73 | output = self.fc(output) |