Update modeling_resnet.py

modeling_resnet.py

@@ -4,26 +4,9 @@ from transformers import PreTrainedModel, PretrainedConfig
 
 class ResNetConfig(PretrainedConfig):
     model_type = "custom_resnet"
-
-    def __init__(
-        self,
-        num_classes=10,
-        image_size=64,
-        input_channels=3,
-        layers=(3, 4, 6, 3),
-        hidden_sizes=(64, 128, 256, 512),
-        activation_function="relu",
-        label_smoothing=0.0,
-        **kwargs
-    ):
+    def __init__(self, num_classes=10, **kwargs):
         super().__init__(**kwargs)
         self.num_classes = num_classes
-        self.image_size = image_size
-        self.input_channels = input_channels
-        self.layers = layers
-        self.hidden_sizes = hidden_sizes
-        self.activation_function = activation_function
-        self.label_smoothing = label_smoothing
 
 class BasicBlock(nn.Module):
     expansion = 1
@@ -51,6 +34,7 @@ class ResNet(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.in_channels = 64
+        # Matches your training script ResNet-34
         self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.relu = nn.ReLU(inplace=True)
@@ -62,26 +46,35 @@ class ResNet(PreTrainedModel):
         self.layer4 = self._make_layer(BasicBlock, 512, 3, stride=2)
         
         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-        self.fc = nn.Linear(512, config.num_classes)
+        self.fc = nn.Linear(512 * BasicBlock.expansion, config.num_classes)
 
     def _make_layer(self, block, out_channels, blocks, stride=1):
         downsample = None
-        if stride != 1 or self.in_channels != out_channels:
+        if stride != 1 or self.in_channels != out_channels * block.expansion:
             downsample = nn.Sequential(
-                nn.Conv2d(self.in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(out_channels),
+                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels * block.expansion),
             )
-        layers = [block(self.in_channels, out_channels, stride, downsample)]
-        self.in_channels = out_channels
+        layers = []
+        layers.append(block(self.in_channels, out_channels, stride, downsample))
+        self.in_channels = out_channels * block.expansion
         for _ in range(1, blocks):
             layers.append(block(self.in_channels, out_channels))
         return nn.Sequential(*layers)
 
-    def forward(self, x):
-        x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
+    def forward(self, x, labels=None):
+        x = self.relu(self.bn1(self.conv1(x)))
+        x = self.maxpool(x)
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)
         x = self.layer4(x)
         x = torch.flatten(self.avgpool(x), 1)
-        return self.fc(x)
+        logits = self.fc(x)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels)
+            
+        return {"loss": loss, "logits": logits} if loss is not labels else logits