diff --git a/train.py b/train.py new file mode 100644 index 0000000..c3ee0ff --- /dev/null +++ b/train.py @@ -0,0 +1,88 @@ +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]) \ No newline at end of file