import torchimport torch.nn as nnimport torch.optim as optimfrom torch.utils.data import DataLoader
# 假设我们有一个自定义的数据集和模型dataset = ...# 你的数据集model = ...# 你的模型，继承自 nn.Moduleloss_function = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)
# 数据加载dataloader = DataLoader(dataset, batch_size=64, shuffle=True)
# 训练循环for epoch in range(num_epochs):for inputs, targets in dataloader:optimizer.zero_grad() # 清零梯度outputs = model(inputs)# 前向传播loss = loss_function(outputs, targets)# 计算损失loss.backward() # 反向传播        optimizer.step()                    # 更新模型参数
print(f"Epoch {epoch}, Loss: {loss.item()}")

 pytorch_lightning  pl

# 单节点多GPU训练trainer = Trainer(gpus=4)
# 多节点分布式训练trainer = Trainer(accelerator="ddp")
# 执行训练前，需要配置训练节点，设置主节点，专用端口# 启动训练，需要指定主节点，如下：python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_PER_NODE  --nnodes=NUM_NODES  --node_rank=NODE_RANK --master_addr=MASTER_ADDR --master_port=MASTER_PORT train.py