File size: 1,505 Bytes
0de77cb |
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 |
from transformers import PreTrainedModel
from .configuration_efficientnetv25 import EfficientNetV25Config
import torch, sys, os
from huggingface_hub import hf_hub_download
class EfficientNetV25ForImageClassification(PreTrainedModel):
config_class = EfficientNetV25Config
def __init__(self, config):
super().__init__(config)
repo_id = '/'.join(config.url.split('/')[3:5])
file_name = config.url.split('/')[-1]
path = f"./models/{file_name}"
if not os.path.exists(path):
hf_hub_download(repo_id=repo_id, filename=file_name, local_dir="./models")
self.model = torch.load(path)
self.input_size = config.input_size
shape = [2] + self.input_size
example_inputs = torch.randn(shape)
example_inputs = (example_inputs - example_inputs.min()) / (example_inputs.max() - example_inputs.min())
self.num_classes = config.num_classes
if self.num_classes != 1000:
self.model.classifier = torch.nn.Linear(in_features=1984, out_features=self.num_classes, bias=True)
traced_model = torch.jit.trace(self.model, example_inputs)
traced_model.save(file_name)
self.model = torch.jit.load(file_name)
def forward(self, tensor, labels=None):
logits = self.model(tensor)
if labels is not None:
loss = torch.nn.cross_entropy(logits, labels)
return {"loss": loss, "logits": logits}
return {"logits": logits} |