Create a Feature extractor from pretrained model and export to ONNX

Hello!

I need a suggestion from the community. I’ve a pretrained classification model (let’s take resnet50 as a reference). I want use such pretrained model as feature extractor, doing some stuff with extracted features, and export it as ONNX.

I’ve tried several approaches, but all of them brutally failed.

Code example:


weights = torchvision.models.ResNet50_Weights.IMAGENET1K_V2
model = torchvision.models.resnet50(weights=weights)
preprocess = weights.transforms()
img = torch.zeros(1, 3, 512, 1024).to("cuda")


class prova(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model.to("cuda")
        self.outputs = ()

        def hook(module, input, output):
            outputs = self.outputs + (output.detach(),)

        self.model.layer1.register_forward_hook(hook)

    def forward(self, x):
        with torch.no_grad():
            print("| embedding extract | test |")
            result = self.model(x)
            embedding_vectors = self.outputs[0]
        return (result,) + (embedding_vectors,)


feat_extractor = prova(model)

feat_extractor(img)

torch.onnx.export(
    feat_extractor,  # --dynamic only compatible with cpu
    img,
    "prova.onnx",
    verbose=False,
    opset_version=13,
    do_constant_folding=True,  # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False
    input_names=["images"],
)

Fail with:
[ CUDAFloatType{1,256,128,256} ]) of traced region did not have observable data dependence with trace inputs; this probably indicates your program cannot be understood by the tracer.

Do you have any suggestion on how to reach the goal? It would be important not modifying the forward pass of the original model as it is embedded in a complex flow.