Hi, I try to create a first-in-first-out queue as a pytorch model. The queue, with a limited size, updates every time when a new input comes, and returns the updated queue. Codes are very simple:
import torch
import torch.nn as nn
class WavBuffer(nn.Module):
def __init__(self, size=10):
super().__init__()
self.size = size
wavbuf = torch.zeros(size)
self.register_buffer('wavbuf', wavbuf)
def forward(self, x):
self.wavbuf = torch.cat([self.wavbuf, x])[-self.size:]
return self.wavbuf
model = WavBuffer(10)
x = torch.ones(5)
for i in range(2):
wavbuf = model(x)
print(wavbuf)
As expected, the outputs are:
tensor([0., 0., 0., 0., 0., 1., 1., 1., 1., 1.])
tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
Then I export the model to onnx format and infer with onnxruntime:
torch.onnx.export(
model, torch.zeros(5), 'model.onnx', verbose=False, input_names=['wav'],
output_names=['wavbuf'], opset_version=11
)
import numpy as np
import onnxruntime
model = onnxruntime.InferenceSession('model.onnx')
x = np.ones(5, dtype=np.float32)
inputs = {model.get_inputs()[0].name: x}
for i in range(2):
outputs = model.run(None, inputs)
wavbuf = outputs[0]
print(wavbuf)
However, now the outputs are:
[0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
[0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
I guess that weights in onnx models are not changeable, but is there any solution to create writable buffers during model design and change the buffers in onnx inference? An available example is LSTM, where the hidden states update for each time step. However, it is too difficult for me to its implementation.