Upload ResNet10

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "ResNet10"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_resnet.ResNet10Config"
+    "AutoConfig": "configuration_resnet.ResNet10Config",
+    "AutoModel": "modeling_resnet.ResNet10"
   },
   "depths": [
     1,
@@ -8,6 +12,7 @@
     1,
     1
   ],
+  "dtype": "float32",
   "embedding_size": 64,
   "hidden_act": "relu",
   "hidden_sizes": [

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:10f7d125770aa256bd45ec9e4f586ca1157e29380fa1306d14a025664ae173d0
-size 19626736
+oid sha256:1be729cd30827ec597a7601d7b10ce1d6a50729b79119142761fbb9bba090d5f
+size 19627000

modeling_resnet.py

@@ -123,6 +123,21 @@ class Conv2dJax(nn.Module):
         return self.conv(x)
 
 
+class MyGroupNorm(nn.Module):
+    def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+        super().__init__()
+        self.group_norm = nn.GroupNorm(num_groups, num_channels, eps, affine)
+
+    def forward(self, x):
+        if x.ndim == 3:
+            x = x.unsqueeze(0)
+            x = self.group_norm(x)
+            x = x.squeeze(0)
+        else:
+            x = self.group_norm(x)
+        return x
+
+
 class BasicBlock(nn.Module):
     def __init__(self, in_channels, out_channels, activation, stride=1, norm_groups=4):
         super().__init__()
@@ -134,17 +149,17 @@ class BasicBlock(nn.Module):
             stride=stride,
             bias=False,
         )
-        self.norm1 = nn.GroupNorm(num_groups=norm_groups, num_channels=out_channels)
+        self.norm1 = MyGroupNorm(num_groups=norm_groups, num_channels=out_channels)
         self.act1 = ACT2FN[activation]
         self.act2 = ACT2FN[activation]
         self.conv2 = Conv2dJax(out_channels, out_channels, kernel_size=3, stride=1, bias=False)
-        self.norm2 = nn.GroupNorm(num_groups=norm_groups, num_channels=out_channels)
+        self.norm2 = MyGroupNorm(num_groups=norm_groups, num_channels=out_channels)
 
         self.shortcut = None
         if in_channels != out_channels:
             self.shortcut = nn.Sequential(
                 Conv2dJax(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
-                nn.GroupNorm(num_groups=norm_groups, num_channels=out_channels),
+                MyGroupNorm(num_groups=norm_groups, num_channels=out_channels),
             )
 
     def forward(self, x):
@@ -190,18 +205,18 @@ class Encoder(nn.Module):
                 )
 
     def forward(self, hidden_state: Tensor, output_hidden_states: bool = False) -> BaseModelOutputWithNoAttention:
-        hidden_states = () if output_hidden_states else None
+        hidden_states: Optional[tuple[Tensor, ...]] = () if output_hidden_states else None
 
         for stage in self.stages:
             if output_hidden_states:
-                hidden_states = hidden_states + (hidden_state,)
+                hidden_states = hidden_states + (hidden_state,)  # type: ignore
 
             hidden_state = stage(hidden_state)
 
         if output_hidden_states:
-            hidden_states = hidden_states + (hidden_state,)
+            hidden_states = hidden_states + (hidden_state,)  # type: ignore
 
-        return BaseModelOutputWithNoAttention(
+        return BaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=hidden_state,
             hidden_states=hidden_states,
         )
@@ -222,47 +237,22 @@ class ResNet10(PreTrainedModel):
                 padding=3,
                 bias=False,
             ),
-            # The original code has a small trick -
-            # https://github.com/rail-berkeley/hil-serl/blob/main/serl_launcher/serl_launcher/vision/resnet_v1.py#L119
-            # class MyGroupNorm(nn.GroupNorm):
-            #     def __call__(self, x):
-            #         if x.ndim == 3:
-            #             x = x[jnp.newaxis]
-            #             x = super().__call__(x)
-            #             return x[0]
-            #         else:
-            #             return super().__call__(x)
-            nn.GroupNorm(num_groups=4, eps=1e-5, num_channels=self.config.embedding_size),
+            MyGroupNorm(num_groups=4, eps=1e-5, num_channels=self.config.embedding_size),
             ACT2FN[self.config.hidden_act],
             MaxPool2dJax(kernel_size=3, stride=2),
         )
 
         self.encoder = Encoder(self.config)
-        self.pooler = nn.AdaptiveAvgPool2d(output_size=1)
+        self._init_pooler()
+        self.post_init()
 
     def _init_pooler(self):
         if self.config.pooler == "avg":
             self.pooler = nn.AdaptiveAvgPool2d(output_size=1)
         elif self.config.pooler == "max":
             self.pooler = nn.MaxPool2d(kernel_size=3, stride=2)
-        elif self.config.pooler == "spatial_learned_embeddings":
-            raise ValueError("Invalid pooler, it exist in the hil serl version, but weights are missing")
-
-            # In the original HIl-SERL code is used SpatialLearnedEmbeddings as pooliing method
-            # Check https://github.com/rail-berkeley/hil-serl/blob/7d17d13560d85abffbd45facec17c4f9189c29c0/serl_launcher/serl_launcher/agents/continuous/sac.py#L490
-            # But weights for this custom layer are missing
-            # Probably it means that pretrained weights used other way of pooling - probably it's AvgPool2d
-            # self.pooler = nn.Sequential(
-            #     SpatialLearnedEmbeddings(
-            #         height=height,
-            #         width=width,
-            #         channel=channel,
-            #         num_features=self.num_spatial_blocks,
-            #     ),
-            #     nn.Dropout(0.1, deterministic=not train),
-            # )
         else:
-            raise ValueError(f"Invalid pooler: {self.config.pooler}")
+            self.pooler = None
 
     def forward(self, x: Tensor, output_hidden_states: Optional[bool] = None) -> BaseModelOutputWithNoAttention:
         output_hidden_states = (
@@ -271,7 +261,10 @@ class ResNet10(PreTrainedModel):
         embedding_output = self.embedder(x)
         encoder_outputs = self.encoder(embedding_output, output_hidden_states=output_hidden_states)
 
-        pooler_output = self.pooler(encoder_outputs.last_hidden_state)
+        if self.pooler is not None:
+            pooler_output = self.pooler(encoder_outputs.last_hidden_state)
+        else:
+            pooler_output = None
 
         return BaseModelOutputWithPoolingAndNoAttention(
             last_hidden_state=encoder_outputs.last_hidden_state,