New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Batch Norm update doesn't require cpu #2

Open

ListIndexOutOfRange opened this issue Jan 6, 2022 · 0 comments

ListIndexOutOfRange commented Jan 6, 2022

Hi,

First, thanks for your work. It's pretty cool !
I just wanted to clarify something.

In your example code, in order to update the batch normalization statistics at the end of training, you wrote:

swa_model = swa_model.cpu()
torch.optim.swa_utils.update_bn(train_loader, swa_model)
swa_model = swa_model.cuda()

but the update_bn function does accept a device keyword argument device in order to use cuda if possible (see https://github.com/pytorch/pytorch/blob/master/torch/optim/swa_utils.py#L124).

So actually a better way to write your code would be something like:

device  = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.optim.swa_utils.update_bn(train_loader, swa_model, device)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment