scaled_mnist.py 文件源码

def train(model, generator, batch_num, epoch):
    model.train()
    for batch_idx in range(batch_num):
        data, target = next(generator)
        data, target = torch.from_numpy(data), torch.from_numpy(target)
        # convert BHWC to BCHW
        data = data.permute(0, 3, 1, 2)
        data, target = data.float().cuda(), target.long().cuda()

        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        loss = F.cross_entropy(output, target)
        loss.backward()
        optimizer.step()

    print('Train Epoch: {}\tLoss: {:.6f}'.format(epoch, loss.data[0]))