File size: 1,170 Bytes
6061356 de4e5d8 6061356 8e66c4a 6061356 12d9fbf 6061356 12d9fbf 6061356 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | import torchvision.models as models
from transformers import PreTrainedModel
import torch.nn as nn
from transformers import PretrainedConfig
class MusheffConfig(PretrainedConfig):
model_type = "efficientnet_b3"
def __init__(self, num_classes=12, dropout_rate=0.3, **kwargs):
self.num_classes = num_classes
self.dropout_rate = dropout_rate
super().__init__(**kwargs)
class Musheff(PreTrainedModel):
config_class = MusheffConfig # Link to config
def __init__(self, config):
super().__init__(config)
# Extract parameters
num_classes = config.num_classes
dropout_rate = config.dropout_rate
# # Load default weights from base model
# weights = models.EfficientNet_B3_Weights.DEFAULT
# Load base model
self.model = models.efficientnet_b3(weights=None)
# Modify classifier head
in_features = self.model.classifier[1].in_features
self.model.classifier = nn.Sequential(
nn.Dropout(p=dropout_rate, inplace=True),
nn.Linear(in_features, num_classes),
)
def forward(self, x):
return self.model(x)
|