一个很典型的train.py

#定义网络
net=Unet()

#设置数据预处理的方式:转为tensor，以及进行归一化
transforms =transforms.Compose([transforms.ToTensor(),
                      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

#训练集数据先下载下来，然后进行预处理
train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
                                    download=True, transform=transform)
#把训练集数据读到内存里面
train_loader = torch.utils.data.DataLoader(train_set, batch_size=4,
                                     shuffle=True, num_workers=2)

#测试集数据先下载下来，然后进行预处理
test_set = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
#把测试集数据读到内存里面
test_loader = torch.utils.data.DataLoader(test_set, batch_size=4,
                                       shuffle=False, num_workers=2)

#定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

#开始训练
model.train()  # 进入训练状态,让你的模型知道现在正在训练
for epoch in range(20):#20个epoch
    loss=0
    for i,data in enumerate(train_loader,0):
        inputs,labels=data
        outputs=model(inputs)
        optimizer.zero_grad()
        loss=criterion(outputs,labels)
        loss.backward()
        optimizer.step()#根据梯度更新参数