用pytorch写一段神经网络

好的，我将用 PyTorch 写一段基本的神经网络代码。

首先，需要导入 PyTorch 库并设定设备（cpu 或 gpu）：

import torch

device = 'cuda' if torch.cuda.is_available() else 'cpu'

然后，定义输入和输出的形状：

input_size = 784  # 输入层的节点数，对应图像大小为 28*28
hidden_size = 500  # 隐藏层的节点数
num_classes = 10  # 输出层的节点数，对应分类数量

接下来，定义神经网络的结构：

class NeuralNet(torch.nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = torch.nn.Linear(input_size, hidden_size)
        self.relu = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(hidden_size, num_classes)

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

model = NeuralNet(input_size, hidden_size, num_classes).to(device)

在上面的代码中，我们定义了一个包含两个全连接层的神经网络，并使用 ReLU 作为激活函数。这个神经网络有一个输入层、一个隐藏层和一个输出层。

然后，我们可以使用 PyTorch 中的各种优化器和损失函数来训练这个神经网络：

learning_rate = 0.001

# 定义损失函数和优化器
criterion = torch.nn.CrossEntropyLoss