I am trying to export the model to onnx.
I can achieve that using the torch.onnx.export()
class net_nn(torch.nn.Module):
def __init__(self, training=True):
....
def load_weights(self, path, args=None):
....
def foward(self, x):
....
model = net_nn(training=False)
model .load_weights(args.trained_model, args=args)
dummy_input = torch.ones(2, 3, dtype=torch.long)
torch.onnx.export(model, dummy_input , 'net_nn.onnx', verbose=True)
But what I want to do further is adding the post-processing in the onnx. Here is what I did,
class net_nn(torch.nn.Module):
def __init__(self, training=True):
....
def load_weights(self, path, args=None):
....
def foward(self, x):
....
model = net_nn(training=False)
model .load_weights(args.trained_model, args=args)
def post_processing(x):
....
class net_and_post(torch.nn.Module):
def foward(self, x):
pred = model(x)
return post_processing(pred)
model_with_post = net_and_post()
dummy_input = torch.ones(2, 3, dtype=torch.long)
torch.onnx.export(model_with_post , dummy_input , 'net_nn_post.onnx', verbose=True)
But it ends by
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
Please help.