在导出onnx模型的时候,输入的observation是ndarray字典,将它输入到一个features_extractors的时候会报下面的错误:
Traceback (most recent call last):
File "E:/Workspace/PybulletVelocityController/Test/testPTh.py", line 48, in <module>
encoded_tensor_list.append(extractor(dummy_input[key]))
File "C:\Users\huyu\anaconda3\envs\drones\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "C:\Users\huyu\anaconda3\envs\drones\lib\site-packages\torch\nn\modules\flatten.py", line 45, in forward
return input.flatten(self.start_dim, self.end_dim)
TypeError: flatten() takes from 0 to 1 positional arguments but 2 were given
debug到函数tensor的flatten函数那块一直觉得没有错误,但就是报错,不知道怎么回事。
后来发现ndarray也有flatten函数。
它的参数刚好是0 或 1。
这就解释了为什么报上面的错,原因是没有调用tensor的flatten方法,而是去调用了ndarray的flatten方法。
在stable-baselines3 里面应该有将ndarray数据预处理成tensor的方法(之前没有注意到),因此在训练过程中没有报这样的数据类型错误。