File size: 744 Bytes
d7abc8f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 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 |