From f525cf3039509ad074ca5397a7a79c249a4a7bba Mon Sep 17 00:00:00 2001
From: suixin1424 <1424531758@qq.com>
Date: Sun, 28 Jan 2024 20:05:30 +0800
Subject: [PATCH] =?UTF-8?q?=E8=AE=AD=E7=BB=83=E5=AE=8C=E8=87=AA=E5=8A=A8?=
 =?UTF-8?q?=E5=AF=BC=E5=87=BAonnx=E6=A8=A1=E5=9E=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 train.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/train.py b/train.py
index c3ee0ff..c4fce2b 100644
--- a/train.py
+++ b/train.py
@@ -1,8 +1,12 @@
+import numpy as np
 import pandas as pd
 import torch
 from torch.utils.data import Dataset, DataLoader
 import torch.nn as nn
 import torch.optim as optim
+import onnxruntime
+import onnx
+from onnxsim import simplify
 
 # 读取CSV文件
 train_csv_path = 'mouse_data.csv'  # 替换为你的CSV文件路径
@@ -85,4 +89,9 @@ for idx, data in enumerate(test_dataloader):
         print(loss)
         if idx == len(test_dataloader)-1:
             print(batch_dx_dy[0])
-            print(output[0])
\ No newline at end of file
+            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'])
\ No newline at end of file