Demo测试类
2024.8.12日基于骨骼点的时序动作多分类项目测试类代码编写
输入到模型的数据是(batchsize,sequence,frames,feature_points)
import pickle
import torch
from torch import nn
test_pkl = "./save_pkl.pkl"
'''
Define the LSTM model
'''
num_joints = 26
num_features_per_joint = 6
num_classes = 18
hidden_size = 128
input_size = num_joints * num_features_per_joint
num_layers = 2
class ActionClassifier_(nn.Module):
def __init__(self):
super(ActionClassifier_, self).__init__()
self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True)
self.fc = nn.Linear(hidden_size * num_layers, num_classes) # *2 for bidirectional
self._export = False
def export(self):
self.eval()
'''
This is used to set whether or not to export onnx filess
'''
self._export = True
'''
Acquisition speed or accelerations
'''
def op_diff(self,t1 :torch.tensor):
b, f, j, c = t1.shape
padding = torch.zeros((b, 1, j, c)).to(device=t1.device)
t2 = t1[:,1:] - t1[:,:-1]
t2 = torch.cat((padding, t2), dim=1)
return t2
def forward(self, x):
b, f, j, c = x.shape
if self._export:
velocity = self.op_diff(x)
acceleration = self.op_diff(velocity)
x_add_a = torch.cat((x, velocity, acceleration), dim=3)
# add_f = torch.zeros((b, f, j, c)).to(device=x.device)
# xy8s = x[:, :, 8, :].unsqueeze(-2)
# xy10s = x[:, :, 10, :].unsqueeze(-2)
# j__ = x[:, :, :, :] - xy8s
# add_f[:, :, :, 0] = j__[:, :, :, 1]/j__[:, :, :, 0]
# j__ = x[:, :, :, :] - xy10s
# add_f[:, :, :, 1] = j__[:, :, :, 1] / j__[:, :, :, 0]
# x_add_f = torch.cat((x_add_a, add_f), dim=3)
x = x_add_a.view(b, f, j *num_features_per_joint)
else:
x = x.view(b, f, j * c)
h0 = torch.zeros(self.lstm.num_layers * 2, b, self.lstm.hidden_size).to(x.device)
c0 = torch.zeros(self.lstm.num_layers * 2, b, self.lstm.hidden_size).to(x.device)
out, _ = self.lstm(x, (h0, c0))
out = self.fc(out)
if self._export:
# sofmax
out = out.softmax(-1)
out = torch.argmax(out,dim=-1)
return out
def export_onnx():
model = ActionClassifier_().to(device=torch.device('cuda'))
model.load_state_dict(torch.load("2024_08_12_17_48_200_90.77.pth"))
with open(test_pkl, "rb") as f:
data = pickle.load(f)
test_data = data["data"]
test_label = data["labels"]
model.export()
y = model(test_data)
# model is model struct and paramter,test_data is input model data, model_whole.onnx is save path,input_names is save onnx file inputname, output_names is save onnx file outputname.
torch.onnx.export(model, test_data , "model_whole.onnx", verbose=True, input_names=['input'], output_names=['output'])
if __name__ == '__main__':
export_onnx()
憧憬
学习使用torch保存onnx模型。