| 6 | |
| 7 | |
| 8 | class ImageEncoder(torch.nn.Module): |
| 9 | def __init__(self, args, keep_lang=False): |
| 10 | super().__init__() |
| 11 | |
| 12 | self.model, self.train_preprocess, self.val_preprocess = clip.load( |
| 13 | args.model, args.device, jit=False |
| 14 | ) |
| 15 | |
| 16 | self.cache_dir = args.cache_dir |
| 17 | |
| 18 | if not keep_lang and hasattr(self.model, "transformer"): |
| 19 | delattr(self.model, "transformer") |
| 20 | |
| 21 | def forward(self, images): |
| 22 | assert self.model is not None |
| 23 | return self.model.encode_image(images) |
| 24 | |
| 25 | def save(self, filename): |
| 26 | print(f"Saving image encoder to {filename}") |
| 27 | utils.torch_save(self, filename) |
| 28 | |
| 29 | def load(self, filename): |
| 30 | print(f"Loading image classifier from {filename}") |
| 31 | return utils.torch_load(self, filename) |
| 32 | |
| 33 | |
| 34 | class ClassificationHead(torch.nn.Linear): |
no outgoing calls
no test coverage detected