| 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} |