Update resnet.py

Adding ResNet 18 model class
GT-RIPL · Feb 17, 2022 · 50d9c2e · 50d9c2e
1 parent b3ffc65
commit 50d9c2e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/models/resnet.py b/models/resnet.py
@@ -112,6 +112,9 @@ def forward(self, x, pen=False):
 def resnet32(out_dim):
     return ResNet(BasicBlock, [5, 5, 5], num_classes=out_dim)
 
+def resnet18(out_dim):
+    return ResNet(BasicBlock, [2, 2, 2, 2], num_classes=out_dim)
+
 class BiasLayer(nn.Module):
     def __init__(self):
         super(BiasLayer, self).__init__()
@@ -122,4 +125,4 @@ def forward(self, x):
         return self.alpha * x + self.beta
 
     def printParam(self, i):
-        print(i, self.alpha.item(), self.beta.item())
+        print(i, self.alpha.item(), self.beta.item())