基于pytorch写一个三层神经网络，训练数据并导出模型

import torch
import torch.nn as nn
import torch.optim as optim

# 定义三层神经网络
class ThreeLayerNN(nn.Module):
    def __init__(self, input_size, hidden_size1, hidden_size2, output_size):
        super(ThreeLayerNN, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size1)
        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
        self.fc3 = nn.Linear(hidden_size2, output_size)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 创建模型实例
input_size = 10  # 输入层大小，根据实际情况调整
hidden_size1 = 32  # 第一层隐藏层大小，根据实际情况调整
hidden_size2 = 16  # 第二层隐藏层大小，根据实际情况调整
output_size = 2  # 输出层大小，根据实际情况调整
model = ThreeLayerNN(input_size, hidden_size1, hidden_size2, output_size)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 根据实际情况选择损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)  # 使用Adam优化器，学习率可调

# 训练数据和标签（此处仅为示例，您需要根据实际情况提供数据）
X_train = torch.randn(100, input_size)  # 随机生成100个样本作为训练数据，输入维度为input_size
Y_train = torch.randint(0, output_size, (100,))  # 随机生成100个标签，输出维度为output_size

# 训练模型
num_epochs = 10  # 训练轮数，可根据实际情况调整
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(X_train)
    loss = criterion(outputs, Y_train)
    
    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    if (epoch + 1) % 1 == 0:
        print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch + 1, num_epochs, loss.item()))

# 导出模型
torch.save(model.state_dict(), 'model.pth')

============