Demo测试类编写

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模型。