[PyTorch] 06. wandb로 학습 진행상황 로깅하기

import wandb

config = {
    'batch_size' : 32,
    'epoch' : 4,
    'some-other-config' : "you can put any value here"
    ...
}

wandb.init(project='my-wandb-project', name='run-1', config=config, save_code=False)

wandb.log({'var1' : value, 'var2' : value2})

def train_model(net, criterion, optimizer, dataloaders_dict, num_epochs, model_name=''):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    net.to(device)
    print(f"Using device: {device}")
    
    ### config값을 저장한 dict 생성 ###
    config = {
        'model_name' : model_name,
        'batch_size' : dataloaders_dict['train'].batch_size,
        'epoch' : num_epochs,
        'criterion' : criterion,
        'optimizer' : optimizer,
        'some-other-config' : 'you can put any value here'
    }
    ### 새 wandb 세션 시작 ###
    wandb.init(reinit=True, project='my-wandb-project', config=config)
    ### 모델 파라미터(그래디언트 등) 추적을 위한 .watch 호출 ###
    wandb.watch(net)
    
    for epoch in range(num_epochs):
        print("-----------------")
        print(f"> Epoch {epoch+1}/{num_epochs}")
        for phase in ['train', 'val']:
            torch.set_grad_enabled(phase == 'train')
            if phase == 'train':
                net.train()
            else:
                net.eval()
            if (epoch == 0) and (phase == 'train'):
                continue
            epoch_loss = 0.0
            epoch_corrects = 0
            for x,y in tqdm(dataloaders_dict[phase]):
                x, y = x.to(device), y.to(device)
                output =  net(x)
                loss = criterion(output, y)
                if phase == 'train':
                    optimizer.zero_grad()
                    loss.backward()
                    optimizer.step()
                _, y_pred = torch.max(output, 1)
                epoch_loss += loss.item() * x.size(0)
                epoch_corrects += torch.sum(y_pred == y.data)
                ### train phase 한정으로 각 batch마다 loss를 기록 ###
                if phase == 'train':
                    wandb.log({'train_loss' : loss})
            epoch_loss = epoch_loss / len(dataloaders_dict[phase].dataset)
            epoch_acc = epoch_corrects.double() / len(dataloaders_dict[phase].dataset)
            ### 각 epoch마다 해당 phase에서의 loss, acc을 기록 ###
            wandb.log({
                f'{phase}_epoch_loss' : epoch_loss,
                f'{phase}_epoch_acc' : epoch_acc
            })
            print(f"{phase} Loss: {epoch_loss:.4f}  Acc: {epoch_acc:.4f} ({epoch_corrects}/{len(dataloaders_dict[phase].dataset)})")
    ### 학습이 끝난 후 세션 종료를 위한 .finish() 호출 ###
    wandb.finish()

# resnet18로 학습 진행
net = models.resnet18()
net.fc = nn.Linear(in_features=net.fc.in_features, out_features=10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(params = net.parameters())
train_model(net, criterion, optimizer, dataloaders_dict, num_epochs=4, model_name='resnet18')

# resnet34로 학습 진행
net = models.resnet34()
net.fc = nn.Linear(in_features=net.fc.in_features, out_features=10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(params = net.parameters())
train_model(net, criterion, optimizer, dataloaders_dict, num_epochs=4, model_name='resnet34')

# vgg16으로 학습 진행
net = models.vgg16(pretrained=True)
net.classifier[6] = nn.Linear(in_features=net.classifier[6].in_features, out_features=10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(params = net.parameters())
train_model(net, criterion, optimizer, dataloaders_dict, num_epochs=4, model_name='vgg16(pretrained)')