kd13
/

vit-nano-patch16-224

Image Classification

custom_vit_nano

Model card Files Files and versions

kd13 commited on 3 days ago

Commit

e6238e3

·

verified ·

1 Parent(s): ba314eb

Update modeling_vit.py

Files changed (1) hide show

modeling_vit.py +4 -4

modeling_vit.py CHANGED Viewed

@@ -3,7 +3,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
-from .configuration_vit import CustomViTNanoV2Config
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
@@ -133,7 +133,7 @@ class Block(nn.Module):
         x = x + self.mlp(self.norm2(x))
         return x
-class CustomViTNanoV2PreTrainedModel(PreTrainedModel):
     config_class = CustomViTNanoV2Config
     base_model_prefix = "custom_vit"
     main_input_name = "pixel_values"
@@ -154,7 +154,7 @@ class CustomViTNanoV2PreTrainedModel(PreTrainedModel):
         elif isinstance(module, RMSNorm):
             nn.init.ones_(module.weight)
-class CustomViTNanoV2ForImageClassification(CustomViTNanoV2PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_classes
@@ -220,4 +220,4 @@ class CustomViTNanoV2ForImageClassification(CustomViTNanoV2PreTrainedModel):
             logits=logits,
         )
-CustomViTNanoV2ForImageClassification.register_for_auto_class("AutoModelForImageClassification")

 import torch.nn.functional as F
 from transformers import PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
+from .configuration_vit import CustomViTNanoConfig
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         x = x + self.mlp(self.norm2(x))
         return x
+class CustomViTNanoPreTrainedModel(PreTrainedModel):
     config_class = CustomViTNanoV2Config
     base_model_prefix = "custom_vit"
     main_input_name = "pixel_values"
         elif isinstance(module, RMSNorm):
             nn.init.ones_(module.weight)
+class CustomViTNanoForImageClassification(CustomViTNanoPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_classes
             logits=logits,
         )
+CustomViTNanoForImageClassification.register_for_auto_class("AutoModelForImageClassification")