Spaces:
Sleeping
Sleeping
import torch | |
import torch.nn as nn | |
class SiameseNetwork(nn.Module): | |
def __init__(self): | |
super(SiameseNetwork, self).__init__() | |
self.cnn1 = nn.Sequential( | |
nn.Conv2d(1, 96, kernel_size=11, stride=1), | |
nn.ReLU(inplace=True), | |
nn.LocalResponseNorm(5, alpha=0.0001, beta=0.75, k=2), | |
nn.MaxPool2d(3, stride=2), | |
nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2), | |
nn.ReLU(inplace=True), | |
nn.LocalResponseNorm(5, alpha=0.0001, beta=0.75, k=2), | |
nn.MaxPool2d(3, stride=2), | |
nn.Dropout2d(p=0.3), | |
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1), | |
nn.ReLU(inplace=True), | |
nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1), | |
nn.ReLU(inplace=True), | |
nn.MaxPool2d(3, stride=2), | |
nn.Dropout2d(p=0.3), | |
) | |
self.fc1 = nn.Sequential( | |
nn.Linear(25600, 1024), | |
nn.ReLU(inplace=True), | |
nn.Dropout2d(p=0.5), | |
nn.Linear(1024, 128), | |
nn.ReLU(inplace=True), | |
nn.Linear(128, 2) | |
) | |
def forward_once(self, x): | |
output = self.cnn1(x) | |
output = output.view(output.size()[0], -1) | |
output = self.fc1(output) | |
return output | |
def forward(self, input1, input2): | |
output1 = self.forward_once(input1) | |
output2 = self.forward_once(input2) | |
return output1, output2 | |
# Function to load the trained model | |
def load_model(model_path): | |
model = SiameseNetwork() | |
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) | |
model.eval() | |
return model | |