Create model.py

model.py

@@ -0,0 +1,185 @@
+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}