Upload model.py

model.py

@@ -0,0 +1,87 @@
+"""
+CIFAR 10
+INPUT - [3, 32, 32]
+"""
+import torch.nn as nn
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, dropout):
+        super(BasicBlock, self).__init__()
+        self.cblock = nn.Sequential(
+            *[
+                self._get_base_layer(
+                    in_channel if i == 0 else out_channel, out_channel, dropout
+                )
+                for i in range(2)
+            ]
+        )
+
+    def _get_base_layer(self, in_channel, out_channel, dropout):
+        return nn.Sequential(
+            nn.Conv2d(
+                in_channel,
+                out_channel,
+                kernel_size=3,
+                padding=1,
+                padding_mode="replicate",
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channel),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return self.cblock(x) + x
+
+
+class DavidPageNet(nn.Module):
+    def __init__(self, channels=[64, 128, 256, 512], dropout=0.01):
+        super(DavidPageNet, self).__init__()
+        self.block0 = self._get_base_layer(3, channels[0], pool=False)
+        self.block1 = nn.Sequential(
+            *[
+                self._get_base_layer(channels[0], channels[1]),
+                BasicBlock(channels[1], channels[1], dropout),
+            ]
+        )
+
+        self.block2 = self._get_base_layer(channels[1], channels[2])
+        self.block3 = nn.Sequential(
+            *[
+                self._get_base_layer(channels[2], channels[3]),
+                BasicBlock(channels[3], channels[3], dropout),
+            ]
+        )
+
+        self.logit = nn.Sequential(
+            nn.MaxPool2d(4),
+            nn.Flatten(),
+            nn.Linear(512, 10),
+        )
+
+    def _get_base_layer(self, in_channel, out_channel, pool=True):
+        return nn.Sequential(
+            nn.Conv2d(
+                in_channel,
+                out_channel,
+                stride=1,
+                padding=1,
+                kernel_size=3,
+                bias=False,
+                padding_mode="replicate",
+            ),
+            nn.MaxPool2d(2) if pool else nn.Identity(),
+            nn.BatchNorm2d(out_channel),
+            nn.ReLU(),
+        )
+
+    def forward(self, x):
+        x = self.block0(x)
+
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+
+        return self.logit(x)