Spaces:
Build error
Build error
| # model_architectures.py | |
| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class ResNet50(nn.Module): | |
| def __init__(self, num_classes=7, channels=3): | |
| super(ResNet50, self).__init__() | |
| # Define layers directly without wrapping in 'resnet' | |
| self.conv_layer_s2_same = nn.Conv2d(channels, 64, kernel_size=7, stride=2, padding=3, bias=False) | |
| self.batch_norm1 = nn.BatchNorm2d(64) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) | |
| # Load pre-trained ResNet50 model | |
| resnet = models.resnet50(pretrained=True) | |
| # Extract layers | |
| self.layer1 = resnet.layer1 | |
| self.layer2 = resnet.layer2 | |
| self.layer3 = resnet.layer3 | |
| self.layer4 = resnet.layer4 | |
| self.avgpool = resnet.avgpool | |
| # Fully connected layers | |
| self.fc1 = nn.Linear(resnet.fc.in_features, num_classes) | |
| # If your model has additional fully connected layers, define them here | |
| # Example: | |
| # self.fc2 = nn.Linear(num_classes, num_classes) | |
| def forward(self, x): | |
| x = self.conv_layer_s2_same(x) | |
| x = self.batch_norm1(x) | |
| x = self.relu(x) | |
| x = self.maxpool(x) | |
| x = self.layer1(x) | |
| x = self.layer2(x) | |
| x = self.layer3(x) | |
| x = self.layer4(x) | |
| x = self.avgpool(x) | |
| x = torch.flatten(x, 1) | |
| x = self.fc1(x) | |
| # If additional fully connected layers are defined, pass x through them | |
| # x = self.fc2(x) | |
| return x | |
| def extract_features(self, x): | |
| x = self.conv_layer_s2_same(x) | |
| x = self.batch_norm1(x) | |
| x = self.relu(x) | |
| x = self.maxpool(x) | |
| x = self.layer1(x) | |
| x = self.layer2(x) | |
| x = self.layer3(x) | |
| x = self.layer4(x) | |
| x = self.avgpool(x) | |
| x = torch.flatten(x, 1) | |
| return x | |
| class LSTMPyTorch(nn.Module): | |
| def __init__(self, input_size, hidden_size, num_layers, num_classes): | |
| super(LSTMPyTorch, self).__init__() | |
| self.hidden_size = hidden_size | |
| # Define separate LSTM layers | |
| self.lstm1 = nn.LSTM(input_size, hidden_size, num_layers=1, batch_first=True) | |
| self.lstm2 = nn.LSTM(hidden_size, hidden_size, num_layers=1, batch_first=True) | |
| self.fc = nn.Linear(hidden_size, num_classes) | |
| def forward(self, x): | |
| h0_1 = torch.zeros(1, x.size(0), self.hidden_size).to(x.device) | |
| c0_1 = torch.zeros(1, x.size(0), self.hidden_size).to(x.device) | |
| out1, _ = self.lstm1(x, (h0_1, c0_1)) | |
| h0_2 = torch.zeros(1, x.size(0), self.hidden_size).to(x.device) | |
| c0_2 = torch.zeros(1, x.size(0), self.hidden_size).to(x.device) | |
| out2, _ = self.lstm2(out1, (h0_2, c0_2)) | |
| out = self.fc(out2[:, -1, :]) | |
| return out | |