TopK Output Shape Warning (PyTorch / Onnx Exported Model)

[iprovalo]

This is a small example of a single linear layer with topK onnx export, producing the indices and values of the topK after applying a linear transformation:

import torch
from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions

device = torch.device('cpu')

class LinearTopkModel(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super(LinearTopkModel, self).__init__()
        self.linear1 = torch.nn.Linear(in_features, out_features)

    def forward(self, x, beam_size, token_index):
        x = self.linear1(x)
        token_tensor = x[token_index]
        out = torch.topk(token_tensor, beam_size)
        return out


def main():
    linear_topk_model = LinearTopkModel(4, 8)

    weights = torch.arange(1., 33.).reshape(8, 4)
    print(weights)

    bias = torch.arange(1., 9.)
    print(bias)

    with torch.no_grad():
        linear_topk_model.linear1.weight.copy_(weights)
        linear_topk_model.linear1.bias.copy_(bias)

    linear_topk_model.eval()

    x = torch.arange(1., 13.)
    x = x.view(1, 3, 4)

    print(x)

    with torch.no_grad():
        print("="*30 + "\n" + str(linear_topk_model(x[0], 2, 0)) +"\n"+ "="*30)
        print("="*30 + "\n" + str(linear_topk_model(x[0], 2, 1)) +"\n"+ "="*30)
        print("="*30 + "\n" + str(linear_topk_model(x[0], 2, 2)) +"\n"+ "="*30)

    path = './temp_out.onnx'

    decoder_output = x[0]
    beam_size = 2
    token_index = 2

    torch.onnx._export(
        linear_topk_model,
        (decoder_output, beam_size, token_index, ),
        path,
        export_params=True,
        opset_version=12,
        input_names=["decoder_output", "beam_size","token_index"],
        output_names=["values","indices"],
        dynamic_axes={
            "decoder_output": {0: "batch", 1: "sequence"}
        }
    )


    options = SessionOptions()
    options.intra_op_num_threads = 1
    options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL

    # Load the model as a graph and prepare the CPU backend
    session = InferenceSession(path, options, providers=['CPUExecutionProvider'])
    session.disable_fallback()

    inputs = {
        "decoder_output": decoder_output.cpu().detach().numpy(),
        "beam_size": [beam_size],
        "token_index": [token_index],
    }
    topk = session.run(None, inputs)
    print(topk)


if __name__ == "__main__":
    main()

Even though this works as is, I am getting a warning for the output shape:

2022-12-27 20:52:46.574967 [W:onnxruntime:, execution_frame.cc:811 VerifyOutputSizes] Expected shape from model of {TopKvalues_dim_0} does not match actual shape of {1,2} for output values

Is there a better way to define the forward output or the dynamic_axes?

Thank you,

IP

[wschin]

If you know the expected output shapes. One hack is to load ONNX and edit it directly.

import onnx
m = onnx.load("temp_out.onnx")
# edit m.graph.output below
...

It's possible to provide dynamic_axis for outputs in torch.onnx._export, but it will be error-prone.

[iprovalo]

Thank you, @wschin !