Pneumonia / model.py
DHEIVER's picture
Create model.py
d7abc8f
"""
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