88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
import pandas as pd
|
|
import torch
|
|
from torch.utils.data import Dataset, DataLoader
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
|
|
# 读取CSV文件
|
|
train_csv_path = 'mouse_data.csv' # 替换为你的CSV文件路径
|
|
train_csv = pd.read_csv(train_csv_path, header=None)
|
|
test_csv_path = 'mouse_data_test.csv' # 替换为你的CSV文件路径
|
|
test_csv = pd.read_csv(test_csv_path, header=None)
|
|
|
|
# 提取dx, dy和标签
|
|
dx_dy_train = train_csv.iloc[:, 0].apply(lambda x: list(map(int, x.split(','))))
|
|
dx_dy_labels_train = train_csv.apply(lambda row: [list(map(int, row[i].split(','))) for i in range(1,10)], axis=1)
|
|
dx_dy_test = test_csv.iloc[:, 0].apply(lambda x: list(map(int, x.split(','))))
|
|
dx_dy_labels_test = test_csv.apply(lambda row: [list(map(int, row[i].split(','))) for i in range(1,10)], axis=1)
|
|
|
|
# 转换为PyTorch Tensor
|
|
dx_dy_train_tensor = torch.Tensor(dx_dy_train.tolist())
|
|
dx_dy_train_tensor = dx_dy_train_tensor.unsqueeze(1)
|
|
labels_train_tensor = torch.Tensor(dx_dy_labels_train.tolist())
|
|
dx_dy_test_tensor = torch.Tensor(dx_dy_test.tolist())
|
|
dx_dy_test_tensor = dx_dy_test_tensor.unsqueeze(1)
|
|
labels_test_tensor = torch.Tensor(dx_dy_labels_test.tolist())
|
|
|
|
# 创建自定义Dataset
|
|
class CustomDataset(Dataset):
|
|
def __init__(self, dx_dy, labels):
|
|
self.dx_dy = dx_dy
|
|
self.labels = labels
|
|
|
|
def __len__(self):
|
|
return len(self.dx_dy)
|
|
|
|
def __getitem__(self, idx):
|
|
return self.dx_dy[idx], self.labels[idx]
|
|
|
|
# 创建Dataset和DataLoader
|
|
train_dataset = CustomDataset(dx_dy_train_tensor, labels_train_tensor)
|
|
test_dataset = CustomDataset(dx_dy_test_tensor, labels_test_tensor)
|
|
batch_size = 64
|
|
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
|
|
test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
|
|
|
|
# 定义神经网络模型
|
|
class SimpleNet(nn.Module):
|
|
def __init__(self):
|
|
super(SimpleNet, self).__init__()
|
|
self.fc1 = nn.Linear(2, 64)
|
|
self.fc2 = nn.Linear(64, 32)
|
|
self.fc3 = nn.Linear(32, 18)
|
|
|
|
def forward(self, input_data):
|
|
x = torch.flatten(input_data, start_dim=1)
|
|
x = self.fc1(x)
|
|
x = nn.ReLU()(x)
|
|
x = self.fc2(x)
|
|
x = nn.ReLU()(x)
|
|
x = self.fc3(x)
|
|
return x.view(-1, 9, 2)
|
|
|
|
# 初始化模型、损失函数和优化器
|
|
model = SimpleNet()
|
|
criterion = nn.MSELoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
|
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.9)
|
|
|
|
# 训练模型
|
|
epochs = 1000
|
|
for epoch in range(epochs):
|
|
for batch_dx_dy, batch_labels in train_dataloader:
|
|
optimizer.zero_grad()
|
|
output = model(batch_dx_dy)
|
|
loss = criterion(output, batch_labels)
|
|
print(loss)
|
|
loss.backward()
|
|
optimizer.step()
|
|
scheduler.step()
|
|
print('test')
|
|
for idx, data in enumerate(test_dataloader):
|
|
batch_dx_dy, batch_labels = data
|
|
output = model(batch_dx_dy)
|
|
loss = criterion(output, batch_labels)
|
|
print(loss)
|
|
if idx == len(test_dataloader)-1:
|
|
print(batch_dx_dy[0])
|
|
print(output[0]) |