| import os | |
| from torch.utils import model_zoo | |
| from classes.fc4.squeezenet.SqueezeNet import SqueezeNet | |
| model_urls = { | |
| 1.0: 'https://download.pytorch.org/models/squeezenet1_0-a815701f.pth', | |
| 1.1: 'https://download.pytorch.org/models/squeezenet1_1-f364aa15.pth', | |
| } | |
| class SqueezeNetLoader: | |
| def __init__(self, version: float = 1.1): | |
| self.__version = version | |
| self.__model = SqueezeNet(self.__version) | |
| def load(self, pretrained: bool = False) -> SqueezeNet: | |
| """ | |
| Returns the specified version of SqueezeNet | |
| @param pretrained: if True, returns a model pre-trained on ImageNet | |
| """ | |
| if pretrained: | |
| path_to_local = os.path.join("assets", "pretrained") | |
| os.environ['TORCH_HOME'] = path_to_local | |
| self.__model.load_state_dict(model_zoo.load_url(model_urls[self.__version])) | |
| return self.__model | |