Hi
I have implemented the Einsum operation and can also successfully converted my Pytorch model to ONNX and to TRT(with Einsum op). only for torch.einsum('nctkv,kvw->nctw')
I set the elements of input[0] to be 1, but the input[0] get in enqueue() are not all 1. I have run the program several times and sometimes they are all 1, but sometimes there are only two 1. I don’t know what is wrong.
for example, I will get one of the following two results:
right input[0]:
1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000
wrong input[0]:
1.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 1.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000 0.000
Can you please assist me in resolving this error?
Environment:
CUDA: 11.0
TRT version: 7.1.3
cudnn:8.0.1
pytorch1.6
this is my code:
You can creat a simple network to replicate this problem as follows
class model(nn.Module):
def __init__(self, in_channel=1):
super().__init__()
self.A = np.ones(shape=(3, 15, 15))
self.A = torch.tensor(self.A, dtype=torch.float32, requires_grad=False)
def forward(self, x: torch.Tensor):
x = x.permute(0, 2, 3, 1, 4).contiguous()
x = torch.einsum('nctkv,kvw->nctw', x, self.A)
return x
input_tensor = torch.ones(size=[1, 3, 1, 1, 15])
Model = model()
Model.cuda()
out = Model(input_tensor)
input_name = ['input']
output_name = ['output']
torch.onnx.export(Model,
input_tensor,
'./gcn.onnx',
input_names=input_name, output_names=output_name,
verbose=True,
opset_version=12
)
and then convert gcn.onnx
to TRT by this Einsum plugin, you will reproduce these results above.