Add PerceiverForImageClassificationConvProcessing

src/transformers/__init__.py

@@ -1139,6 +1139,7 @@
             "PerceiverClassificationDecoder",
             "PerceiverForImageClassification",
             "PerceiverForImageClassificationFourier",
+            "PerceiverForImageClassificationConvProcessing",
             "PerceiverForMaskedLM",
             "PerceiverImagePreprocessor",
             "PerceiverLayer",
@@ -3010,6 +3011,7 @@
             PerceiverClassificationDecoder,
             PerceiverForImageClassification,
             PerceiverForImageClassificationFourier,
+            PerceiverForImageClassificationConvProcessing,
             PerceiverForMaskedLM,
             PerceiverImagePreprocessor,
             PerceiverLayer,

src/transformers/models/perceiver/__init__.py

@@ -35,6 +35,7 @@
         "PerceiverClassificationDecoder",
         "PerceiverForImageClassification",
         "PerceiverForImageClassificationFourier",
+        "PerceiverForImageClassificationConvProcessing",
         "PerceiverForMaskedLM",
         "PerceiverLayer",
         "PerceiverModel",
@@ -62,6 +63,7 @@
             PerceiverClassificationDecoder,
             PerceiverForImageClassification,
             PerceiverForImageClassificationFourier,
+            PerceiverForImageClassificationConvProcessing,
             PerceiverForMaskedLM,
             PerceiverLayer,
             PerceiverModel,

src/transformers/models/perceiver/convert_perceiver_haiku_to_pytorch.py

@@ -32,6 +32,7 @@
     PerceiverFeatureExtractor,
     PerceiverForImageClassification,
     PerceiverForImageClassificationFourier,
+    PerceiverForImageClassificationConvProcessing,
     PerceiverForMaskedLM,
     PerceiverTokenizer,
 )
@@ -61,7 +62,7 @@ def rename_keys(state_dict):
                 "trainable_position_encoding/pos_embs", "input_preprocessor.position_embeddings.weight"
             )
 
-        # rename image preprocessor embeddings (for image classification model)
+        # rename image preprocessor embeddings (for image classification model with learned position embeddings)
         name = name.replace("image_preprocessor/~/conv2_d/w", "input_preprocessor.convnet_1x1.weight")
         name = name.replace("image_preprocessor/~/conv2_d/b", "input_preprocessor.convnet_1x1.bias")
         name = name.replace(
@@ -77,6 +78,15 @@ def rename_keys(state_dict):
             "input_preprocessor.positions_projection.bias",
         )
 
+        # rename image preprocessor embeddings (for image classification model with conv processing)
+        if "counter" in name or "hidden" in name:
+            continue
+        name = name.replace("image_preprocessor/~/conv2_d_downsample/~/conv/w", "input_preprocessor.convnet.conv.weight")
+        name = name.replace("image_preprocessor/~/conv2_d_downsample/~/batchnorm/offset", "input_preprocessor.convnet.batchnorm.bias")
+        name = name.replace("image_preprocessor/~/conv2_d_downsample/~/batchnorm/scale", "input_preprocessor.convnet.batchnorm.weight")
+        name = name.replace("image_preprocessor/~/conv2_d_downsample/~/batchnorm/~/mean_ema/average", "input_preprocessor.convnet.batchnorm.running_mean")
+        name = name.replace("image_preprocessor/~/conv2_d_downsample/~/batchnorm/~/var_ema/average", "input_preprocessor.convnet.batchnorm.running_var")
+
         ## DECODERS ##
 
         # rename prefix of decoders
@@ -167,6 +177,10 @@ def rename_keys(state_dict):
         if name[-6:] == "weight" and "embeddings" not in name:
             param = np.transpose(param)
 
+        # if batchnorm, we need to squeeze it
+        if "batchnorm" in name:
+            param = np.squeeze(param)
+
         state_dict["perceiver." + name] = torch.from_numpy(param)
 
 
@@ -180,10 +194,11 @@ def convert_perceiver_checkpoint(pickle_file, pytorch_dump_folder_path, task="ML
     with open(pickle_file, "rb") as f:
         checkpoint = pickle.loads(f.read())
 
+    state = None
     if isinstance(checkpoint, dict):
-        if task not in ["image_classification", "image_classification_fourier"]:
+        if task not in ["image_classification", "image_classification_fourier", "image_classification_conv"]:
             raise ValueError("Make sure to set task to image classification")
-        # the image classification checkpoint with conv_preprocessing also has batchnorm state
+        # the image classification_conv checkpoint also has batchnorm states (running_mean and running_var)
         params = checkpoint["params"]
         state = checkpoint["state"]
     else:
@@ -195,6 +210,12 @@ def convert_perceiver_checkpoint(pickle_file, pytorch_dump_folder_path, task="ML
         for param_name, param in parameters.items():
             state_dict[scope_name + "/" + param_name] = param
 
+    if state is not None:
+        # add state variables
+        for scope_name, parameters in hk.data_structures.to_mutable_dict(state).items():
+            for param_name, param in parameters.items():
+                state_dict[scope_name + "/" + param_name] = param
+
     # rename keys
     rename_keys(state_dict)
 
@@ -226,6 +247,9 @@ def convert_perceiver_checkpoint(pickle_file, pytorch_dump_folder_path, task="ML
         elif task == "image_classification_fourier":
             config.d_model = 261
             model = PerceiverForImageClassificationFourier(config)
+        elif task == "image_classification_conv":
+            config.d_model = 322
+            model = PerceiverForImageClassificationConvProcessing(config)
         else:
             raise ValueError(f"Task {task} not supported")
     else:
@@ -244,7 +268,7 @@ def convert_perceiver_checkpoint(pickle_file, pytorch_dump_folder_path, task="ML
         encoding.input_ids[0, 51:60] = tokenizer.mask_token_id
         inputs = encoding.input_ids
         input_mask = encoding.attention_mask
-    elif task in ["image_classification", "image_classification_fourier"]:
+    elif task in ["image_classification", "image_classification_fourier", "image_classification_conv"]:
         feature_extractor = PerceiverFeatureExtractor()
         image = prepare_img()
         encoding = feature_extractor(image, return_tensors="pt")
@@ -272,7 +296,7 @@ def convert_perceiver_checkpoint(pickle_file, pytorch_dump_folder_path, task="ML
         print("Predicted string:")
         print(tokenizer.decode(masked_tokens_predictions))
 
-    elif task in ["image_classification", "image_classification_fourier"]:
+    elif task in ["image_classification", "image_classification_fourier", "image_classification_conv"]:
         print("Predicted class:", model.config.id2label[logits.argmax(-1).item()])
 
     # Finally, save files

src/transformers/models/perceiver/modeling_perceiver.py

@@ -998,6 +998,80 @@ def forward(
         )
 
 
+# @add_start_docstrings("""Example use of Perceiver for image classification. """, PERCEIVER_START_DOCSTRING)
+class PerceiverForImageClassificationConvProcessing(PerceiverPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.perceiver = PerceiverModel(
+            config,
+            input_preprocessor=PerceiverImagePreprocessor(
+                config,
+                prep_type="conv",
+                spatial_downsample=1,
+                position_encoding_type="fourier",
+                concat_pos=True,
+                max_resolution=(56, 56),
+                num_bands=64,
+                sine_only=False,
+            ),
+            decoder=PerceiverClassificationDecoder(config, num_channels=config.d_latents, use_query_residual=True),
+        )
+
+        self.init_weights()
+
+    # @add_start_docstrings_to_model_forward(PERCEIVER_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    # @add_code_sample_docstrings(
+    #     tokenizer_class=_TOKENIZER_FOR_DOC,
+    #     checkpoint=_CHECKPOINT_FOR_DOC,
+    #     output_type=SequenceClassifierOutput,
+    #     config_class=_CONFIG_FOR_DOC,
+    # )
+    def forward(
+        self,
+        inputs=None,
+        attention_mask=None,
+        head_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        labels=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.perceiver(
+            inputs=inputs,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            if self.num_labels == 1:
+                #  We are doing regression
+                loss_fct = MSELoss()
+                loss = loss_fct(logits.view(-1), labels.view(-1))
+            else:
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
 class PerceiverAbstractDecoder(nn.Module, metaclass=abc.ABCMeta):
     """Perceiver abstract decoder."""
 

src/transformers/models/perceiver/processing_perceiver.py

@@ -17,12 +17,69 @@
 """
 import abc
 from typing import Optional
+import math
+from functools import reduce
+from operator import __add__
 
 import numpy as np
 import torch
 import torch.nn as nn
 
 
+#  ------------------------------------------------------------
+#  -------------------  Up/down-sampling  ---------------------
+#  ------------------------------------------------------------
+
+
+class Conv2dSamePadding(nn.Conv2d):
+    """Conv2d layer with padding="same" support.
+    Source: https://gist.github.com/sumanmichael/4de9dee93f972d47c80c4ade8e149ea6
+    """
+
+    def __init__(self,*args,**kwargs):
+        super(Conv2dSamePadding, self).__init__(*args, **kwargs)
+        self.zero_pad_2d = nn.ZeroPad2d(reduce(__add__,
+            [(k // 2 + (k - 2 * (k // 2)) - 1, k // 2) for k in self.kernel_size[::-1]]))
+
+    def forward(self, input):
+        return  self._conv_forward(self.zero_pad_2d(input), self.weight, self.bias)
+
+
+class Conv2DDownsample(nn.Module):
+  """Downsamples 4x by applying a 2D convolution and doing max pooling."""
+
+  def __init__(
+      self,
+      num_layers: int = 1,
+      in_channels: int = 3,
+      out_channels: int = 64,
+      use_batchnorm: bool = True,
+  ):
+    """Constructs a Conv2DDownsample model.
+    Args:
+      in_channels: The number of input channels.
+      out_channels: The number of conv output channels.
+      use_batchnorm: Whether to use batchnorm.
+    """
+    super().__init__()
+
+    self.conv = Conv2dSamePadding(in_channels=in_channels, out_channels=out_channels,
+          kernel_size=7,
+          stride=2,
+          bias=False)
+    self.batchnorm = nn.BatchNorm2d(num_features=out_channels) if use_batchnorm else nn.Identity()
+    self.relu = nn.ReLU()
+    self.max_pool = nn.MaxPool2d(kernel_size=3, stride=2)
+
+  def forward(self, inputs):
+    out = inputs
+    out = self.conv(inputs)
+    out = self.batchnorm(out)
+    out = self.relu(out)
+    out = self.max_pool(out)
+    return out
+
+
 def generate_fourier_features(pos, num_bands, max_resolution=(224, 224), concat_pos=True, sine_only=False):
     """
     Generate a Fourier frequency position encoding with linear spacing.
@@ -240,18 +297,32 @@ def __init__(
         self.conv_after_patching = conv_after_patching
         self.out_channels = out_channels
 
-        if self.prep_type in ["conv", "patches"]:
-            raise NotImplementedError(f"Preparation type {prep_type} is not yet supported")
+        if self.prep_type == 'conv':
+            # Downsampling with conv is currently restricted
+            convnet_num_layers = math.log(spatial_downsample, 4)
+            convnet_num_layers_is_int = (convnet_num_layers == np.round(convnet_num_layers))
+            if not convnet_num_layers_is_int or temporal_downsample != 1:
+                raise ValueError('Only powers of 4 expected for spatial '
+                                'and 1 expected for temporal '
+                                'downsampling with conv.')
+            self.convnet = Conv2DDownsample(
+                num_layers=int(convnet_num_layers),
+                out_channels=out_channels,
+                use_batchnorm=conv2d_use_batchnorm)
+
         elif self.prep_type == "conv1x1":
-            assert temporal_downsample == 1, "conv1x1 does not downsample in time."
+            if temporal_downsample != 1:
+                raise ValueError("Conv1x1 does not downsample in time.")
             self.convnet_1x1 = nn.Conv2d(
                 in_channels=in_channels,
                 out_channels=out_channels,
                 kernel_size=[1, 1],
                 # spatial_downsample is unconstrained for 1x1 convolutions.
                 stride=(spatial_downsample, spatial_downsample),
             )
-
+        elif self.prep_type == "patches":
+            raise NotImplementedError(f"Preparation type {prep_type} is not yet supported")
+
         if self.position_encoding_type == "trainable":
             self.position_embeddings = PerceiverTrainablePositionEncoding(**position_encoding_kwargs)
         elif self.position_encoding_type == "fourier":
@@ -329,8 +400,13 @@ def _build_network_inputs(self, inputs: torch.Tensor, pos: torch.Tensor, network
         return inputs_with_pos, inputs
 
     def forward(self, inputs: torch.Tensor, pos: Optional[torch.Tensor] = None, network_input_is_1d: bool = True):
-        if self.prep_type in ["conv", "patches"]:
+        if self.prep_type == "patches":
             raise NotImplementedError("TODO")
+        elif self.prep_type == "conv":
+            # Convnet image featurization.
+            # Downsamples spatially by a factor of 4
+            inputs = self.convnet(inputs)
+
         elif self.prep_type == "conv1x1":
             # map inputs to self.out_channels
             inputs = self.convnet_1x1(inputs)