Create model.py

model.py

@@ -0,0 +1,27 @@
+"""
+Builds Pytorch model
+"""
+import torch
+import torchvision.models
+from torch import nn
+
+
+class ResNet101(nn.Module):
+    """
+    ResNet101 model specified for the binary problem. The according transforms were taken from pytorch.org.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.weights = torchvision.models.ResNet101_Weights.DEFAULT
+        self.transforms = self.weights.transforms
+        self.resnet = torchvision.models.resnet101(weights=self.weights)
+
+        for param in self.resnet.parameters():
+            param.requires_grad = False
+
+        self.resnet.fc = nn.Linear(in_features=2048, out_features=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.resnet(x)
+        return x