| import torch
|
| import torch.nn as nn
|
|
|
| from src.logger import get_logger
|
|
|
| logger = get_logger(__name__)
|
|
|
|
|
| class ResidualBlock(nn.Module):
|
| """
|
| Basic ResNet18 Block
|
|
|
| ┌───────────────┐
|
| │ Shortcut │
|
| └──────┬────────┘
|
| │
|
| Conv3x3 -> BN -> ReLU
|
| │
|
| Conv3x3 -> BN
|
| │
|
| Add Shortcut
|
| │
|
| ReLU
|
| """
|
|
|
| def __init__(
|
| self,
|
| in_channels,
|
| out_channels,
|
| stride=1
|
| ):
|
| super().__init__()
|
|
|
| self.conv1 = nn.Conv2d(
|
| in_channels,
|
| out_channels,
|
| kernel_size=3,
|
| stride=stride,
|
| padding=1,
|
| bias=False
|
| )
|
|
|
| self.bn1 = nn.BatchNorm2d(
|
| out_channels
|
| )
|
|
|
| self.relu = nn.ReLU(
|
| inplace=True
|
| )
|
|
|
| self.conv2 = nn.Conv2d(
|
| out_channels,
|
| out_channels,
|
| kernel_size=3,
|
| stride=1,
|
| padding=1,
|
| bias=False
|
| )
|
|
|
| self.bn2 = nn.BatchNorm2d(
|
| out_channels
|
| )
|
|
|
| self.shortcut = nn.Sequential()
|
|
|
|
|
| if (
|
| stride != 1
|
| or
|
| in_channels != out_channels
|
| ):
|
| self.shortcut = nn.Sequential(
|
| nn.Conv2d(
|
| in_channels,
|
| out_channels,
|
| kernel_size=1,
|
| stride=stride,
|
| bias=False
|
| ),
|
| nn.BatchNorm2d(
|
| out_channels
|
| )
|
| )
|
|
|
| def forward(
|
| self,
|
| x
|
| ):
|
|
|
| identity = self.shortcut(
|
| x
|
| )
|
|
|
| out = self.conv1(
|
| x
|
| )
|
|
|
| out = self.bn1(
|
| out
|
| )
|
|
|
| out = self.relu(
|
| out
|
| )
|
|
|
| out = self.conv2(
|
| out
|
| )
|
|
|
| out = self.bn2(
|
| out
|
| )
|
|
|
| out += identity
|
|
|
| out = self.relu(
|
| out
|
| )
|
|
|
| return out
|
|
|
|
|
| class ResNet18(nn.Module):
|
|
|
| def __init__(
|
| self,
|
| num_classes=11
|
| ):
|
| super().__init__()
|
|
|
| logger.info(
|
| "Building ResNet18"
|
| )
|
|
|
|
|
|
|
| self.conv1 = nn.Conv2d(
|
| in_channels=3,
|
| out_channels=64,
|
| kernel_size=3,
|
| stride=1,
|
| padding=1,
|
| bias=False
|
| )
|
|
|
| self.bn1 = nn.BatchNorm2d(
|
| 64
|
| )
|
|
|
| self.relu = nn.ReLU(
|
| inplace=True
|
| )
|
|
|
|
|
|
|
| self.layer1 = nn.Sequential(
|
| ResidualBlock(
|
| 64,
|
| 64
|
| ),
|
| ResidualBlock(
|
| 64,
|
| 64
|
| )
|
| )
|
|
|
|
|
|
|
| self.layer2 = nn.Sequential(
|
| ResidualBlock(
|
| 64,
|
| 128,
|
| stride=2
|
| ),
|
| ResidualBlock(
|
| 128,
|
| 128
|
| )
|
| )
|
|
|
|
|
|
|
| self.layer3 = nn.Sequential(
|
| ResidualBlock(
|
| 128,
|
| 256,
|
| stride=2
|
| ),
|
| ResidualBlock(
|
| 256,
|
| 256
|
| )
|
| )
|
|
|
|
|
|
|
| self.layer4 = nn.Sequential(
|
| ResidualBlock(
|
| 256,
|
| 512,
|
| stride=2
|
| ),
|
| ResidualBlock(
|
| 512,
|
| 512
|
| )
|
| )
|
|
|
| self.global_pool = (
|
| nn.AdaptiveAvgPool2d(
|
| (1, 1)
|
| )
|
| )
|
|
|
| self.dropout = nn.Dropout(
|
| 0.5
|
| )
|
|
|
| self.fc = nn.Linear(
|
| 512,
|
| num_classes
|
| )
|
|
|
| logger.info(
|
| "ResNet18 created successfully"
|
| )
|
|
|
| def forward(
|
| self,
|
| x
|
| ):
|
|
|
| x = self.conv1(
|
| x
|
| )
|
|
|
| x = self.bn1(
|
| x
|
| )
|
|
|
| x = self.relu(
|
| x
|
| )
|
|
|
| x = self.layer1(
|
| x
|
| )
|
|
|
| x = self.layer2(
|
| x
|
| )
|
|
|
| x = self.layer3(
|
| x
|
| )
|
|
|
| x = self.layer4(
|
| x
|
| )
|
|
|
| x = self.global_pool(
|
| x
|
| )
|
|
|
| x = torch.flatten(
|
| x,
|
| 1
|
| )
|
|
|
| x = self.dropout(
|
| x
|
| )
|
|
|
| x = self.fc(
|
| x
|
| )
|
|
|
| return x
|
|
|
|
|
| def build_model():
|
|
|
| model = ResNet18(
|
| num_classes=11
|
| )
|
|
|
| return model
|
|
|
|
|
| if __name__ == "__main__":
|
|
|
| model = build_model()
|
|
|
| print(model)
|
|
|
| sample = torch.randn(
|
| 8,
|
| 3,
|
| 32,
|
| 32
|
| )
|
|
|
| output = model(
|
| sample
|
| )
|
|
|
| print(
|
| "\nOutput Shape:",
|
| output.shape
|
| ) |
