学习用torch写模型

用torch写一个topk+add的小模型

import torch

print(torch.__version__)

class Net1(torch.nn.Module):
    def __init__(self):
        super(Net1, self).__init__()

    def forward(self, x):
        output = torch.topk(x,1000)
        print("**************",type(output),"********")
        input_add = torch.randint(1,10000,(1,1000)).to(torch.int64)
        output_1 = torch.add(input=output[0],other=input_add)


        return output_1

model = Net1()

dummy_input = torch.randint(1, 10000,(1,10336))
dummy_input = dummy_input.to(torch.int64)
input_names = ["input_1"]
# output_names = ["output_1","output_2"]
output_names = ["output_1"]

torch.onnx.export(model, dummy_input, "topk.onnx",input_names=input_names,output_names=output_names,opset_version=12)