making pretrained class too

app.py

@@ -7,6 +7,7 @@ import numpy as np
 from PIL import Image
 from scipy import special
 import sys
+import timm
 from types import SimpleNamespace
 from transformers import AutoModel, pipeline
 import torch
@@ -14,7 +15,7 @@ import torch
 sys.path.insert(1, "../")
 # from utils import model_utils, train_utils, data_utils, run_utils
 # from model_utils import jason_regnet_maker, jason_efficientnet_maker
-from model_utils.efficientnet_config import EfficientNetConfig
+from model_utils.efficientnet_config import EfficientNetConfig, EfficientNetPreTrained
 
 model_path = 'chlab/'
 # model_path = './models/'
@@ -221,10 +222,17 @@ def predict_and_analyze(model_name, num_channels, dim, input_channel, image):
     
     config.save_pretrained(model_loading_name)
     
+    config = EfficientNetConfig.from_pretrained(model_loading_name)
+    
+    model = EfficientNetPreTrained(config)
+
+    pretrained_model = timm.create_model(model_loading_name, pretrained=True)
+    model.model.load_state_dict(pretrained_model.state_dict())
+    
     # pipeline = pipeline(task="image-classification", model=model_loading_name)
         
     # model = load_model(model_name, activation=True)
-    model = AutoModel.from_pretrained(model_loading_name)
+    # model = AutoModel.from_pretrained(model_loading_name)
     print("Model loaded")
     
     print("Looking at activations")

model_utils/efficientnet_config.py

@@ -1,4 +1,4 @@
-from transformers import PretrainedConfig
+from transformers import PretrainedConfig, PreTrainedModel
 from typing import List
 
 import copy
@@ -271,140 +271,51 @@ class EfficientNetConfig(PretrainedConfig):
         """
         
         
-        self.model = EfficientNet(
-                                    dropout=dropout,
-                                    num_channels=num_channels,
-                                    num_classes=num_classes,
-                                    size=size,
-                                    stochastic_depth_prob=stochastic_depth_prob,
-                                    width_mult=width_mult,
-                                    depth_mult=depth_mult,
-        )
+        # self.model = EfficientNet(
+        #                             dropout=dropout,
+        #                             num_channels=num_channels,
+        #                             num_classes=num_classes,
+        #                             size=size,
+        #                             stochastic_depth_prob=stochastic_depth_prob,
+        #                             width_mult=width_mult,
+        #                             depth_mult=depth_mult,
+        # )
         
+        # 
+        self.dropout=dropout
+        self.num_channels=num_channels
+        self.num_classes=num_classes
+        self.size=size
+        self.stochastic_depth_prob=stochastic_depth_prob
+        self.width_mult=width_mult
+        self.depth_mult=depth_mult
+                                    
         super().__init__(**kwargs)
         
-        # super().__init__()
-        # _log_api_usage_once(self)
-        # inverted_residual_setting, last_channel = _efficientnet_conf(
-        #              "efficientnet_%s" % (size), width_mult=width_mult, depth_mult=depth_mult
-        #             )
-
-        # if not inverted_residual_setting:
-        #     raise ValueError("The inverted_residual_setting should not be empty")
-        # elif not (
-        #     isinstance(inverted_residual_setting, Sequence)
-        #     and all([isinstance(s, _MBConvConfig) for s in inverted_residual_setting])
-        # ):
-        #     raise TypeError(
-        #         "The inverted_residual_setting should be List[MBConvConfig]"
-        #     )
-
-        # if "block" in kwargs:
-        #     warnings.warn(
-        #         "The parameter 'block' is deprecated since 0.13 and will be removed 0.15. "
-        #         "Please pass this information on 'MBConvConfig.block' instead."
-        #     )
-        #     if kwargs["block"] is not None:
-        #         for s in inverted_residual_setting:
-        #             if isinstance(s, MBConvConfig):
-        #                 s.block = kwargs["block"]
-
-        # if norm_layer is None:
-        #     norm_layer = nn.BatchNorm2d
-
-        # layers: List[nn.Module] = []
-
-        # # building first layer
-        # firstconv_output_channels = inverted_residual_setting[0].input_channels
-        # layers.append(
-        #     Conv2dNormActivation(
-        #         num_channels,
-        #         firstconv_output_channels,
-        #         kernel_size=3,
-        #         stride=2,
-        #         norm_layer=norm_layer,
-        #         activation_layer=nn.SiLU,
-        #     )
-        # )
-
-        # # building inverted residual blocks
-        # total_stage_blocks = sum(cnf.num_layers for cnf in inverted_residual_setting)
-        # stage_block_id = 0
-        # for cnf in inverted_residual_setting:
-        #     stage: List[nn.Module] = []
-        #     for _ in range(cnf.num_layers):
-        #         # copy to avoid modifications. shallow copy is enough
-        #         block_cnf = copy.copy(cnf)
-
-        #         # overwrite info if not the first conv in the stage
-        #         if stage:
-        #             block_cnf.input_channels = block_cnf.out_channels
-        #             block_cnf.stride = 1
-
-        #         # adjust stochastic depth probability based on the depth of the stage block
-        #         sd_prob = (
-        #             stochastic_depth_prob * float(stage_block_id) / total_stage_blocks
-        #         )
-
-        #         stage.append(block_cnf.block(block_cnf, sd_prob, norm_layer))
-        #         stage_block_id += 1
-
-        #     layers.append(nn.Sequential(*stage))
-
-        # # building last several layers
-        # lastconv_input_channels = inverted_residual_setting[-1].out_channels
-        # lastconv_output_channels = (
-        #     last_channel if last_channel is not None else 4 * lastconv_input_channels
-        # )
-        # layers.append(
-        #     Conv2dNormActivation(
-        #         lastconv_input_channels,
-        #         lastconv_output_channels,
-        #         kernel_size=1,
-        #         norm_layer=norm_layer,
-        #         activation_layer=nn.SiLU,
-        #     )
-        # )
-
-        # self.features = nn.Sequential(*layers)
-        # self.avgpool = nn.AdaptiveAvgPool2d(1)
-        # self.classifier = nn.Sequential(
-        #     nn.Dropout(p=dropout, inplace=True),
-        #     nn.Linear(lastconv_output_channels, num_classes),
-        # )
-
-        # for m in self.modules():
-    #         if isinstance(m, nn.Conv2d):
-    #             nn.init.kaiming_normal_(m.weight, mode="fan_out")
-    #             if m.bias is not None:
-    #                 nn.init.zeros_(m.bias)
-    #         elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
-    #             nn.init.ones_(m.weight)
-    #             nn.init.zeros_(m.bias)
-    #         elif isinstance(m, nn.Linear):
-    #             init_range = 1.0 / math.sqrt(m.out_features)
-    #             nn.init.uniform_(m.weight, -init_range, init_range)
-    #             nn.init.zeros_(m.bias)
-                
-    #     super().__init__(**kwargs)
-
-    # def _forward_impl(self, x: Tensor) -> Tensor:
-    #     x = self.features(x)
-
-    #     x = self.avgpool(x)
-    #     x = torch.flatten(x, 1)
-
-    #     x = self.classifier(x)
-
-    #     return x
-
-    # def forward(self, x: Tensor) -> Tensor:
-    #     return self._forward_impl(x)
+    
+class EfficientNetPreTrained(PreTrainedModel):
+    
+    config_class = EfficientNetConfig
+    
+    def __init__(
+        self,
+        config
+    )
+        super().__init__(config)   
+        self.model = EfficientNet(  dropout=config.dropout,
+                                    num_channels=config.num_channels,
+                                    num_classes=config.num_classes,
+                                    size=config.size,
+                                    stochastic_depth_prob=config.stochastic_depth_prob,
+                                    width_mult=config.width_mult,
+                                    depth_mult=config.depth_mult,)
+        
+    def forward(self, tensor):
+        return self.model.forward(tensor)
     
     
 class EfficientNet(nn.Module):
     
-    model_type = "efficientnet"
     
     def __init__(
         self,