jia-zhuang / pytorch-multi-gpu-training Goto Github PK

# Convert BatchNorm to SyncBatchNorm. 
net = nn.SyncBatchNorm.convert_sync_batchnorm(net)

for epoch in range(10):
    acc_num=0
    for i, (inputs, labels) in enumerate(train_loader):
        # forward
        inputs = inputs.to(device)
        labels = labels.to(device)
        outputs = model(inputs)
        loss = criterion(outputs[0], labels)
        # backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # log
        if args.local_rank == 0 and i % 5 == 0:
            tb_writer.add_scalar('loss', loss.item(), i)
        acc_num += (outputs[0].argmax(1)==labels).sum()
    if args.local_rank == 0:
        tb_writer.add_scalar('acc', acc_num/len(train_dataset),epoch)
        print(f"acc:{acc_num/len(train_dataset)}")