elun15/image-regression

1import torch2import torchvision3 4resnets = {5    "resnet18": torchvision.models.resnet18,6    "resnet34": torchvision.models.resnet34,7    "resnet50": torchvision.models.resnet50,8    "resnet101": torchvision.models.resnet101,9    "resnet152": torchvision.models.resnet152,10}11 12 13# Define a Custom ResNet Model that inherits from torch.nn.Module and modify the Linear layer14class CustomResNet(torch.nn.Module):15    def __init__(self, cfg):16        super(CustomResNet, self).__init__()17 18        # Load the pretrained ResNet-18 model19        resnet = resnets[cfg.model_name](pretrained=True)20        # Output features of the pretrained model21        out_features = resnet.fc.in_features22 23        # Define the Backbone as all layers except the last Linear layer24        self.model = torch.nn.Sequential(*list(resnet.children())[:-1])25 26        # Modify the last Linear layer to have 1 output and the same input as the pretrained model27        # self.fc = torch.nn.Sequential(28        #     torch.nn.Linear(out_features, 256),29        #     torch.nn.BatchNorm1d(256),30        #     torch.nn.ReLU(),31        #     torch.nn.Linear(256, 64),32        #     torch.nn.BatchNorm1d(64),33        #     torch.nn.ReLU(),34        #     torch.nn.Linear(64, 1),35        #     # sigmoid36        #     torch.nn.Sigmoid(),37        # )38        self.fc = torch.nn.Sequential(39            torch.nn.Linear(out_features, 256),40            torch.nn.BatchNorm1d(256),41            torch.nn.ReLU(),42            torch.nn.Linear(256, 1),43            torch.nn.Sigmoid(),44        )45 46    def forward(self, x):47        # Backbone48        x = self.model(x)49        # Flatten50        x = x.view(x.size(0), -1)51        # Head52        x = self.fc(x)53        return x54