File size: 2,031 Bytes
4b7c478 | 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 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import torch
import torchvision
import torchvision.models as models
import torch.nn as nn
def get_model_transfer_learning(model_name="resnet18", n_classes=50):
# Get the requested architecture
if hasattr(models, model_name):
model_transfer = getattr(models, model_name)(pretrained=True)
else:
torchvision_major_minor = ".".join(torchvision.__version__.split(".")[:2])
raise ValueError(f"Model {model_name} is not known. List of available models: "
f"https://pytorch.org/vision/{torchvision_major_minor}/models.html")
# Freeze all parameters in the model
# HINT: loop over all parameters. If "param" is one parameter,
# "param.requires_grad = False" freezes it
# MY CODE HERE
for param in model_transfer.parameters():
param.requires_grad = False
# Add the linear layer at the end with the appropriate number of classes
# 1. get numbers of features extracted by the backbone
num_ftrs = 512
# 2. Create a new linear layer with the appropriate number of inputs and
# outputs
model_transfer.fc = nn.Linear(num_ftrs,n_classes)
return model_transfer
######################################################################################
# TESTS
######################################################################################
import pytest
@pytest.fixture(scope="session")
def data_loaders():
from .data import get_data_loaders
return get_data_loaders(batch_size=2)
def test_get_model_transfer_learning(data_loaders):
model = get_model_transfer_learning(n_classes=23)
dataiter = iter(data_loaders["train"])
images, labels = next(dataiter)
out = model(images)
assert isinstance(
out, torch.Tensor
), "The output of the .forward method should be a Tensor of size ([batch_size], [n_classes])"
assert out.shape == torch.Size(
[2, 23]
), f"Expected an output tensor of size (2, 23), got {out.shape}"
|