train.py 文件源码

def train():
    densenet.train()
    corrects = total_loss = 0
    for data, label in tqdm(training_data, mininterval=1,
                desc='Train Processing', leave=False):
        data, label = Variable(data), Variable(label)
        if use_cuda:
            data, label = data.cuda(), label.cuda()

        optimizer.zero_grad()

        target = densenet(data)
        loss = criterion(target, label)

        loss.backward()
        optimizer.step()

        total_loss += loss.data
        corrects += (torch.max(target, 1)[1].view(label.size()).data == label.data).sum()

    return total_loss[0]/training_size, corrects, corrects/training_size * 100.0

# ##############################################################################
# Save Model
# ##############################################################################