[Pytorch] model.zero_grad() vs. optimizer.zero_grad()

gradient accumulation 관련 코드를 찾아보다 model.zero_grad()를 처음 보게 되어서 글 짧게 남깁니다.

 

def train(opt):
	optimizer = torch.optim.Adam(model.parameters(), ...)
    
	for epoch in opt.n_epochs:
    	model.train()
        for step, batch in enumerate(train_loader):
        	optimizer.zero_grad()
        	
            x, y = batch[0].to(opt.device), batch[1].to(device)
            output = model(x)
            
            loss = criterion(output, y)
            
            loss.backward() # backward 시 gradient가 계산되며 텐서의 grad 속성에 저장됨.
            optimizer.step()

 

위의 optimizer는 model의 파라미터를 받아서 해당 파라미터 텐서만을 optimize하게 됩니다. 보통, 스텝마다 zero_grad()를 해주지 않으면 이전 스텝에서의 gradient가 누적되어서 현재 step에서 optimizer가 최적화를 진행할 때에 의도하지 않은 방향으로 최적화가 진행될 수 있기 때문에 스텝마다 zero_grad()를 해주게 됩니다.

여기서의 zero_grad()는 optimizer에 등록된 파라미터의 gradient를 0으로 다시 초기화해주는 과정입니다. 반면, model.zero_grad()는 해당 모델 내에 있는 모든 파라미터 텐서의 gradient를 0으로 초기화해주게 됩니다. 따라서, optimizer에 해당 model의 모든 파라미터를 등록해두었다면 model.zero_grad()나 optimizer.zero_grad() 나 동일하게 기능하게 될 것입니다.

하지만, optimizer에 모델의 파라미터 중 일부분만 등록을 해두거나, 두 개 이상의 모델 인스턴스의 파라미터를 등록해둔다면 model.zero_grad()와 optimizer.zero_grad()는 다르게 기능할 것입니다. 이런 부분을 염두에 두고 zero_grad()를 해주면 되겠습니다!

상기한 내용에 오류나 의문점이 있을 경우 댓글 부탁드립니다!