File size: 2,687 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import torch
import torch.nn as nn
# define the CNN architecture
# MY CODE HERE
class ResidualBlock(nn.Module):
def __init__(self, in_ch, out_ch, stride=1):
super().__init__()
self.conv1 = nn.Conv2d(in_ch, out_ch, 3, stride, 1, bias=False)
self.bn1 = nn.BatchNorm2d(out_ch)
self.relu = nn.ReLU(True)
self.conv2 = nn.Conv2d(out_ch, out_ch, 3, 1, 1, bias=False)
self.bn2 = nn.BatchNorm2d(out_ch)
self.shortcut = nn.Identity()
if stride != 1 or in_ch != out_ch:
self.shortcut = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 1, stride, bias=False),
nn.BatchNorm2d(out_ch)
)
def forward(self, x):
out = self.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
out += self.shortcut(x)
return self.relu(out)
class MyModel(nn.Module):
def __init__(self, num_classes: int = 1000, dropout: float = 0.7) -> None:
super().__init__()
self.Backbone = nn.Sequential(
ResidualBlock(3, 64, 1),
ResidualBlock(64, 128, 2),
ResidualBlock(128, 256, 2),
ResidualBlock(256, 512, 2),
nn.AdaptiveAvgPool2d((1,1))
)
self.Classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(512, 512),
nn.BatchNorm1d(512),
nn.ReLU(True),
nn.Dropout(dropout),
nn.Linear(512,192),
nn.BatchNorm1d(192),
nn.ReLU(True),
nn.Dropout(dropout),
nn.Linear(192, num_classes)
)
def forward(self, x):
return self.Classifier(self.Backbone(x))
######################################################################################
# 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_model_construction(data_loaders):
model = MyModel(num_classes=23, dropout=0.3)
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}"
|