Just want to check my code with experts

1

Hi guys. I’m trying to do the classification task.
and want to check my code with you guys

  1. is there any part is wrong with pytorch grammar matter
  2. and is there better way to write code? any suggestion
def save_model(model, optimizer, scheduler,idx) :
    model_cpu = model.to('cpu')
    state = {
        'model' : model_cpu.state_dict(),
        'optimizer' : optimizer.state_dict(),
        'scheduler' : scheduler.state_dict()
    }
    if not(os.path.isdir('./saved_model')) : os.mkdir('./saved_model')
    torch.save(state, './saved_model/saved_model.pth')


def train():
    model.train().cuda()
    print('==============================train===============================')
    lst_out = []
    lst_label=[]
    for idx2,(images,labels) in enumerate(train_loader):
        images,labels = torch.tensor(images),torch.tensor(labels)
        images = images.cuda()
        labels = labels.cuda()
        outputs = model(images).cuda()
        loss = criterion(outputs, labels)
        if idx2%args.iter_num ==0: print('train->{}, loss->{}'.format((idx2)/(len(train_loader)),loss))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()



def valid():
    print('==============================valid===============================')
    torch.cuda.empty_cache()
    lst_out = []
    lst_label=[]
    with torch.no_grad():
        model.eval().cuda()
        for idx,( image, label) in enumerate(valid_loader):
            image = image.cuda()
            pred = model(image)
            lst_out += torch.max(pred,1).indices
            label = label.cuda()
            lst_label += label
        f1_scr = f1_score(torch.tensor(lst_out), torch.tensor(lst_label), average='macro')
    return f1_scr



def test() :
    print('==============================test===============================')
    torch.cuda.empty_cache()
    test_loader = data_loader(phase='test', batch_size= args.ts_bs)
    os.system('rm -rf ./test_result.csv')
    f1 = open('./test_result.csv', 'w', encoding ='utf-8-sig', newline='')
    wr = csv.writer(f1)
    lst_out = []
    lst_label = []
    ct = 0
    model = MODEL
    model.load_state_dict(torch.load('./saved_model/saved_model.pth')['model'])
    with torch.no_grad():
        torch.eval().cuda()
        for idx,( image, _) in enumerate(test_loader):
            image = image.cuda()
            pred = model(image)
            pred = pred.argmax(1).cpu().detach().numpy().astype('int')
            lst_out.extend(pred)
            if idx%iter_num:print('idx[',idx,']/',len(test_loader),len(lst_out))
        for idx, item in enumerate(lst_out) :
            wr.writerows([[dic_label[item]]])


if __name__ == "__main__":

    f1_scr = 0
    model = MODEL
    model.cuda()
    criterion = nn.CrossEntropyLoss()

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.001)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size = 5, gamma =0.1)
    best_acc,loss = 0,0

    train_loader = data_loader(phase='train', batch_size=args.bs)
    valid_loader = data_loader(phase='validation',batch_size=8)
    print('main')

    for i in range(args.epoch):
        train()
        f1_result = valid()
        if f1_result > f1_scr:
            f1_scr = f1_result
            save_model(model,optimizer,scheduler,idx)
    #if args.status =='train':train()
2

Your code looks generally fine beside some small things I would change.

  • outputs = model(images).cuda(): the cuda() call is not necessary, since the output should already be on the GPU (or are you transferring it to the CPU inside the model?) Also, to write device-agnostic code, you could use:
device = 'cuda:0'`
images = images.to(device)
labels = labels.to(device)
...

  • torch.cuda.empty_cache(): Clearing the cache is usually not necessary and will slow down your code. You might want to use it, if you expect other applications to use GPU memory and thus want PyTorch to release it.

  • model.eval().cuda(): unnecessary cuda() operation

  • NIT: lst_out += torch.max(pred,1).indices vs. pred = pred.argmax(1).cpu().detach().numpy().astype('int'): you are getting the predictions in both lines of code, but use quite a different coding style. Not an issue at all, so just a minor remark.

  • model = MODEL: this looks a bit fishy and I guess MODEL might just be a placeholder for this post?

  • NIT: step_size = 5, gamma =0.1 and best_acc,loss = 0,0: your spacing is inconsistent. You could use a linter to clean it up.

  • train() and valid() method use global variables. Personally I would pass all variables as arguments as I’ve ran into issues, where I’ve changed the global variable name in the script and the method was still using an “old” variable.

1 Like