Upload folder using huggingface_hub

configuration_resnet.py

@@ -1,7 +1,5 @@
 from transformers import PretrainedConfig
 from typing import List
-from pprint import pprint
-
 
 class ResnetConfig(PretrainedConfig):
     model_type = "faen_resnet"
@@ -34,13 +32,3 @@ class ResnetConfig(PretrainedConfig):
         self.stem_type = stem_type
         self.avg_down = avg_down
         super().__init__(**kwargs)
-
-if __name__ == "__main__":
-    resnet50d_config = ResnetConfig(block_type="bottleneck", stem_width=32, stem_type="deep", avg_down=True)
-    print("init a ResnetConfig, it is:\n")
-    pprint(resnet50d_config)
-    resnet50d_config.save_pretrained("./")
-    resnet50d_config = ResnetConfig.from_pretrained("./")
-    print("\n")
-    print("saved to file config.json and reload it from config.json and it is:\n")
-    pprint(resnet50d_config)

modeling_resnet.py

@@ -1,12 +1,13 @@
 from transformers import PreTrainedModel
 from timm.models.resnet import BasicBlock, Bottleneck, ResNet
-from configuration_resnet import ResnetConfig
+from .configuration_resnet import ResnetConfig
 import torch
 
 BLOCK_MAPPING = {"basic": BasicBlock, "bottleneck": Bottleneck}
 
 
 class ResnetModel(PreTrainedModel):
+    config_class = ResnetConfig
 
     def __init__(self, config):
         super().__init__(config)
@@ -28,7 +29,8 @@ class ResnetModel(PreTrainedModel):
     
 
 class ResnetModelForImageClassification(PreTrainedModel):
-    
+    config_class = ResnetConfig
+
     def __init__(self, config):
         super().__init__(config)
         block_layer = BLOCK_MAPPING[config.block_type]