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)[0][1]
probs
tensor(0.0722, grad_fn=<SelectBackward>)
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?