How to train and infer models in PyTorch?

In PyTorch, training and inference of models typically involve the following steps:

1. Defining model: To begin, the structure of the neural network model needs to be defined, which can be done by creating a custom neural network model through inheritance of the torch.nn.Module class.
2. Define the loss function: Choose an appropriate loss function to calculate the difference between the model’s predicted values and the true labels.
3. Definition of optimizer: Selecting the appropriate optimizer to update model parameters, commonly used optimizers include SGD, Adam, etc.
4. Train the model: During the training process, input data is passed into the model through iterative loops, calculating loss and updating parameters until a stopping condition is reached.

model = YourModel()  # 定义模型
criterion = torch.nn.CrossEntropyLoss()  # 定义损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)  # 定义优化器

for epoch in range(num_epochs):
    for inputs, labels in train_loader:
        optimizer.zero_grad()  # 清空梯度
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

# 推理
model.eval()  # 切换到评估模式
with torch.no_grad():
    for inputs, labels in test_loader:
        outputs = model(inputs)
        # 进行推理操作

During the training process, additional features can be added as needed, such as learning rate adjustment strategies, model saving and loading, etc. Finally, during the inference stage, the model should be switched to evaluation mode and the gradients should be turned off using the torch.no_grad() context manager to speed up inference.