I have a re-trained resne50 model, the model structure is the following
model = torch.load('./model/best_finetuning_resnet50_.pth') device = torch.device('cpu') model = model.to(device) model.eval() model
ResNet( (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) (relu): ReLU(inplace) (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) (layer1): Sequential( ..... ..... (avgpool): AvgPool2d(kernel_size=7, stride=1, padding=0) (fc): Linear(in_features=2048, out_features=2, bias=True) )
The model output is a two dimensional tensor
dummy_inputs = torch.rand(1, 3, 224, 224) outputs = model(dummy_inputs) outputs
tensor([[ 1.2956, -1.2583]], grad_fn=<ThAddmmBackward>)
However, the final result I need is just the second element, i.e., the probability the input is labeled as one
probs = torch.softmax(outputs, dim = 1) probs
I would like to make ‘probs’ to be an end-to-end model such that the output is 1 dimension probability value while the input is 3x224x224. I need to output the model as ONNX so that it can be used by the third party who can simply feed the batch_size x 3 x 224 x 224 dimension input into the model to get the batch_size x 1 dimension probability values.
torch.onnx.export(probs, dummy_inputs, "./model/best_finetuning_resnet50_.onnx", verbose=True, input_names='input_feature')
The above code does not work as this just saved a tensor Not a model. How can I make this work?