yellowdog/mouse_control
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

训练完自动导出onnx模型

Browse Source
This commit is contained in:
suixin1424
2024-01-28 20:05:30 +08:00
committed by GitHub
parent 3ab0d75e6e
commit f525cf3039
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
train.py
View File

@@ -1,8 +1,12 @@
import numpy as np
import pandas as pd import pandas as pd
import torch import torch
from torch.utils.data import Dataset, DataLoader from torch.utils.data import Dataset, DataLoader
import torch.nn as nn import torch.nn as nn
import torch.optim as optim import torch.optim as optim
import onnxruntime
import onnx
from onnxsim import simplify
# 读取CSV文件 # 读取CSV文件
train_csv_path = 'mouse_data.csv' # 替换为你的CSV文件路径 train_csv_path = 'mouse_data.csv' # 替换为你的CSV文件路径
@@ -86,3 +90,8 @@ for idx, data in enumerate(test_dataloader):
if idx == len(test_dataloader)-1: if idx == len(test_dataloader)-1:
print(batch_dx_dy[0]) print(batch_dx_dy[0])
print(output[0]) print(output[0])
model.eval()
onnx_name = 'mouse.onnx'
dummy = torch.randn(1, 1, 2)
torch.onnx.export(model, dummy, onnx_name,verbose=True, input_names=['input'], output_names=['output'])