| 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): |