|
|
|
@ -51,12 +51,13 @@ def torch_no_grad(f): |
|
|
|
# @accelerate |
|
|
|
class Model: |
|
|
|
def __init__(self, model_name, device, num_classes): |
|
|
|
self.device = device |
|
|
|
self.model = create_model(model_name, pretrained=True, num_classes=num_classes) |
|
|
|
self.model.eval() |
|
|
|
self.model.to(device) |
|
|
|
|
|
|
|
def __call__(self, x: torch.Tensor): |
|
|
|
return self.model.forward_features(x) |
|
|
|
return self.model.forward_features(x.to(self.device)) |
|
|
|
|
|
|
|
|
|
|
|
@register(output_schema=['vec']) |
|
|
|
@ -113,7 +114,7 @@ class TimmImage(NNOperator): |
|
|
|
img = img if self.skip_tfms else self.tfms(img) |
|
|
|
img_list.append(img) |
|
|
|
inputs = torch.stack(img_list) |
|
|
|
inputs = inputs.to(self.device) |
|
|
|
inputs = inputs |
|
|
|
features = self.model(inputs) |
|
|
|
if isinstance(features, list): |
|
|
|
features = [self.post_proc(x) for x in features] |
|
|
|
@ -143,6 +144,7 @@ class TimmImage(NNOperator): |
|
|
|
return img |
|
|
|
|
|
|
|
def post_proc(self, features): |
|
|
|
features = features.to(self.device) |
|
|
|
if features.dim() == 3: |
|
|
|
features = features[:, 0] |
|
|
|
if features.dim() == 4: |
|
|
|
@ -166,7 +168,7 @@ class TimmImage(NNOperator): |
|
|
|
path = path + '.onnx' |
|
|
|
else: |
|
|
|
raise AttributeError(f'Invalid format {format}.') |
|
|
|
dummy_input = torch.rand((1,) + self.config['input_size']).to(self.device) |
|
|
|
dummy_input = torch.rand((1,) + self.config['input_size']) |
|
|
|
if format == 'pytorch': |
|
|
|
torch.save(self._model, path) |
|
|
|
elif format == 'torchscript': |
|
|
|
@ -182,7 +184,7 @@ class TimmImage(NNOperator): |
|
|
|
elif format == 'onnx': |
|
|
|
self._model.forward = self._model.forward_features |
|
|
|
try: |
|
|
|
torch.onnx.export(self._model, |
|
|
|
torch.onnx.export(self._model.to('cpu'), |
|
|
|
dummy_input, |
|
|
|
path, |
|
|
|
input_names=['input_0'], |
|
|
|
|
