Upload model

hf_model.py

@@ -20,6 +20,7 @@ from transformers import PretrainedConfig, PreTrainedModel
 
 
 from .model import create_model_from_args
+from .model import RADIOModel as RADIOModelBase
 from .input_conditioner import get_default_conditioner, InputConditioner
 
 
@@ -42,7 +43,11 @@ class RADIOConfig(PretrainedConfig):
 
 
 class RADIOModel(PreTrainedModel):
-    """Pretrained Hugging Face model for RADIO."""
+    """Pretrained Hugging Face model for RADIO.
+
+    This classes inherits from both PreTrainedModel, which provides
+    HuggingFace's functionality for loading and saving models.
+    """
 
     config_class = RADIOConfig
 
@@ -52,32 +57,19 @@ class RADIOModel(PreTrainedModel):
         RADIOArgs = namedtuple("RADIOArgs", config.args.keys())
         args = RADIOArgs(**config.args)
         self.config = config
-        self.model = create_model_from_args(args)
-        self.input_conditioner: InputConditioner = get_default_conditioner()
-
-    def forward(self, x: torch.Tensor):
-        x = self.input_conditioner(x)
+        model = create_model_from_args(args)
+        input_conditioner: InputConditioner = get_default_conditioner()
 
-        y = self.model.forward_features(x)
+        self.radio_model = RADIOModelBase(
+            model,
+            input_conditioner,
+            config.return_summary,
+            config.return_spatial_features,
+        )
 
-        if isinstance(y, (list, tuple)):
-            summary, all_feat = y
-        elif isinstance(self.model, VisionTransformer):
-            patch_gen = getattr(self.model, "patch_generator", None)
-            if patch_gen is not None:
-                summary = y[:, : patch_gen.num_cls_tokens].flatten(1)
-                all_feat = y[:, patch_gen.num_skip :]
-            elif self.model.global_pool == "avg":
-                summary = y[:, self.model.num_prefix_tokens :].mean(dim=1)
-                all_feat = y
-            else:
-                summary = y[:, 0]
-                all_feat = y[:, 1:]
-        else:
-            raise ValueError("Unsupported model type")
+    @property
+    def model(self) -> VisionTransformer:
+        return self.radio_model.model
 
-        if self.config.return_summary and self.config.return_spatial_features:
-            return summary, all_feat
-        elif self.config.return_summary:
-            return summary
-        return all_feat
+    def forward(self, x: torch.Tensor):
+        return self.radio_model.forward(x)

model.py

@@ -6,11 +6,56 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
+import torch
 from torch import nn
 
-from timm.models import create_model
+from timm.models import create_model, VisionTransformer
 
 from .enable_cpe_support import enable_cpe
+from .input_conditioner import InputConditioner
+
+
+class RADIOModel(nn.Module):
+    def __init__(
+        self,
+        model: nn.Module,
+        input_conditioner: InputConditioner,
+        return_summary: bool,
+        return_spatial_features: bool,
+    ):
+        super().__init__()
+
+        self.model = model
+        self.input_conditioner = input_conditioner
+        self.return_summary = return_summary
+        self.return_spatial_features = return_spatial_features
+
+    def forward(self, x: torch.Tensor):
+        x = self.input_conditioner(x)
+
+        y = self.model.forward_features(x)
+
+        if isinstance(y, (list, tuple)):
+            summary, all_feat = y
+        elif isinstance(self.model, VisionTransformer):
+            patch_gen = getattr(self.model, "patch_generator", None)
+            if patch_gen is not None:
+                summary = y[:, : patch_gen.num_cls_tokens].flatten(1)
+                all_feat = y[:, patch_gen.num_skip :]
+            elif self.model.global_pool == "avg":
+                summary = y[:, self.model.num_prefix_tokens :].mean(dim=1)
+                all_feat = y
+            else:
+                summary = y[:, 0]
+                all_feat = y[:, 1:]
+        else:
+            raise ValueError("Unsupported model type")
+
+        if self.return_summary and self.return_spatial_features:
+            return summary, all_feat
+        elif self.return_summary:
+            return summary
+        return all_feat
 
 
 def create_model_from_args(args) -> nn.Module:
@@ -36,13 +81,16 @@ def create_model_from_args(args) -> nn.Module:
         **args.model_kwargs,
     )
 
-    assert not args.cls_token_per_teacher or args.cpe_max_size is not None, "CPE must be enabled for multiple CLS tokens!"
+    assert (
+        not args.cls_token_per_teacher or args.cpe_max_size is not None
+    ), "CPE must be enabled for multiple CLS tokens!"
 
     if args.cpe_max_size is not None:
-        enable_cpe(model,
-                   args.cpe_max_size,
-                   num_cls_tokens=len(args.teachers) if args.cls_token_per_teacher else 1,
-                   register_multiple=args.register_multiple,
+        enable_cpe(
+            model,
+            args.cpe_max_size,
+            num_cls_tokens=len(args.teachers) if args.cls_token_per_teacher else 1,
+            register_multiple=args.register_multiple,
         )
 
     return model

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:242360b04b7f78204b535ce8a96e28ef3316520d55be43e6873fd45696fb9d61
-size 2662619441
+oid sha256:ad369b92359d9a42f93f6bbb9be2191f79b4b6fc923fdd31d992ca32336f608d
+size 2662624177