| from huggingface_hub import PyTorchModelHubMixin |
| from torch import nn |
|
|
| class SurfinBird(nn.Module, PyTorchModelHubMixin): |
| def __init__(self, config: dict) -> None: |
| super().__init__() |
| self.conv1 = nn.Conv2d( |
| in_channels=config["num_channels"], |
| out_channels=64, |
| kernel_size=7, |
| stride=2, |
| padding=3) |
| self.bn1 = nn.BatchNorm2d(64) |
| self.relu1 = nn.ReLU() |
| self.mp1 = nn.MaxPool2d(kernel_size=2, |
| stride=2) |
| self.conv_block_2 = nn.Sequential( |
| nn.Conv2d( |
| in_channels=64, |
| out_channels=64, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(64), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=64, |
| out_channels=64, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(64), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=64, |
| out_channels=64, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(64), |
| nn.ReLU(), |
| nn.MaxPool2d(kernel_size=2, |
| stride=2) |
| ) |
| self.conv_block_3 = nn.Sequential( |
| nn.Conv2d( |
| in_channels=64, |
| out_channels=128, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(128), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=128, |
| out_channels=128, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(128), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=128, |
| out_channels=128, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(128), |
| nn.ReLU(), |
| nn.MaxPool2d(kernel_size=2, |
| stride=2) |
| ) |
| self.conv_block_4 = nn.Sequential( |
| nn.Conv2d( |
| in_channels=128, |
| out_channels=128, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(128), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=128, |
| out_channels=128, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(128), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=128, |
| out_channels=128, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(128), |
| nn.ReLU(), |
| nn.MaxPool2d(kernel_size=2, |
| stride=2) |
| ) |
| self.conv_block_5 = nn.Sequential( |
| nn.Conv2d( |
| in_channels=128, |
| out_channels=256, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(256), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=256, |
| out_channels=256, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(256), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=256, |
| out_channels=256, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(256), |
| nn.ReLU(), |
| nn.MaxPool2d(kernel_size=2, |
| stride=2) |
| ) |
| self.conv_block_6 = nn.Sequential( |
| nn.Conv2d( |
| in_channels=256, |
| out_channels=256, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(256), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=256, |
| out_channels=256, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(256), |
| nn.ReLU(), |
| nn.Conv2d( |
| in_channels=256, |
| out_channels=256, |
| kernel_size=3, |
| stride=1, |
| padding=1 |
| ), |
| nn.BatchNorm2d(256), |
| nn.ReLU(), |
| nn.MaxPool2d(kernel_size=2, |
| stride=2) |
| ) |
|
|
| self.avgpool = nn.Sequential( |
| nn.AdaptiveAvgPool2d(output_size=(1, 1)) |
| ) |
|
|
| self.classifier = nn.Sequential( |
| nn.Flatten(), |
| nn.Linear(in_features=config["hidden_units"]*1*1, |
| out_features=config["num_classes"]) |
| ) |
| def forward(self, x: torch.Tensor): |
| return self.classifier(self.avgpool(self.conv_block_6(self.conv_block_5(self.conv_block_4(self.conv_block_3(self.conv_block_2(self.mp1(self.relu1(self.bn1(self.conv1(x))))))))))) |
|
|
| config = {"num_channels": 3, "hidden_units": 256, "num_classes": 525} |
