[FIXED] Which opset version is supported by torch.oonnx

Issue

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.

Solution

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

Answered By - TheProbPro