Update modeling_resnet.py

modeling_resnet.py

@@ -4,9 +4,10 @@ from transformers import PreTrainedModel, PretrainedConfig
 
 class ResNetConfig(PretrainedConfig):
     model_type = "custom_resnet"
-    def __init__(self, num_classes=10, **kwargs):
+    def __init__(self, num_classes=10, num_channels=3, **kwargs):
         super().__init__(**kwargs)
         self.num_classes = num_classes
+        self.num_channels = num_channels  
 
 class BasicBlock(nn.Module):
     expansion = 1
@@ -34,8 +35,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.conv1 = nn.Conv2d(config.num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
         self.bn1 = nn.BatchNorm2d(64)
         self.relu = nn.ReLU(inplace=True)
         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
@@ -46,18 +46,17 @@ class ResNet(PreTrainedModel):
         self.layer4 = self._make_layer(BasicBlock, 512, 3, stride=2)
         
         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
-        self.fc = nn.Linear(512 * BasicBlock.expansion, config.num_classes)
+        self.fc = nn.Linear(512, config.num_classes)
 
     def _make_layer(self, block, out_channels, blocks, stride=1):
         downsample = None
-        if stride != 1 or self.in_channels != out_channels * block.expansion:
+        if stride != 1 or self.in_channels != out_channels:
             downsample = nn.Sequential(
-                nn.Conv2d(self.in_channels, out_channels * block.expansion, kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(out_channels * block.expansion),
+                nn.Conv2d(self.in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels),
             )
-        layers = []
-        layers.append(block(self.in_channels, out_channels, stride, downsample))
-        self.in_channels = out_channels * block.expansion
+        layers = [block(self.in_channels, out_channels, stride, downsample)]
+        self.in_channels = out_channels
         for _ in range(1, blocks):
             layers.append(block(self.in_channels, out_channels))
         return nn.Sequential(*layers)
@@ -72,9 +71,4 @@ class ResNet(PreTrainedModel):
         x = torch.flatten(self.avgpool(x), 1)
         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
+        return {"logits": logits}