Im trying to export the SPIGA model to .onnx format. My code to do so is the following:
# Create three dummy tensors with the specified sizes
input_image = torch.randn(1, 3, 256, 256).cuda() # Size: (batch_size, channels, height, width)
landmarks = torch.randn(1, 98, 3).cuda() # Size: (batch_size, num_landmarks, 3)
cam_matrix = torch.randn(1, 3, 3).cuda() # Size: (batch_size, 3, 3)
# Create a list containing these tensors
dummy_input = [input_image, landmarks, cam_matrix]
#dummy_input = ([torch.randn(1, 3, 256, 256).cuda(), self.model3d, self.cam_matrix])
onnx_model_path = "spiga_model.onnx" # Output ONNX file path
torch.onnx.export(
self.model, # Y/our SPIGA model instance
dummy_input, # Example input data
onnx_model_path, # Output ONNX file path
verbose=True, # Enable verbose mode for debugging (optional)
input_names=self.model_inputs, # List of input names (adjust as needed)
output_names=["features"], # List of output names (adjust as needed)
opset_version=14 # ONNX opset version (adjust as needed)
)
print(f"SPIGA model exported to {onnx_model_path}")
From what i could read from the documentation the opset version should be between 9 and 16. I have tried versio 7 to 20 and it tells me that none of these are supported.
Im using pycharm and this is my pytorch version: Version: 2.0.1+cu117
and my torchvision version is: Version: 0.15.2+cu117.
The error code im getting from pycharm is: Exporting the operator 'aten::affine_grid_generator' to ONNX opset version 14 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues. its the same error message I get with all the opset versions i have tried.
The error was with the torch.nn.functional.affine_grid operator. Therefore I implemented the following code found from this post: https://github.com/pytorch/pytorch/issues/30563
the code used is:
def affine_grid(theta, size, align_corners=False):
N, C, H, W = size
grid = create_grid(N, C, H, W)
grid = grid.view(N, H * W, 3).bmm(theta.transpose(1, 2))
grid = grid.view(N, H, W, 2)
return grid
def create_grid(N, C, H, W):
grid = torch.empty((N, H, W, 3), dtype=torch.float32).cuda()
grid.select(-1, 0).copy_(linspace_from_neg_one(W))
grid.select(-1, 1).copy_(linspace_from_neg_one(H).unsqueeze_(-1))
grid.select(-1, 2).fill_(1)
return grid
def linspace_from_neg_one(num_steps, dtype=torch.float32):
r = torch.linspace(-1, 1, num_steps, dtype=torch.float32)
r = r * (num_steps - 1) / num_steps
return r
This fixed my issue and let me export my model to .onnx format