The to_share_or_not_to_share from apourchot

to_share_or_not_to_share's People

Contributors

Stargazers

Watchers

to_share_or_not_to_share's Issues

Forget to write back the batchnorm statistics while finetuning the bn layer

I really like the idea that you should finetune the bn layers before evaluating the model. I find you try to finetune the bn layers of the sampled model in function ft_bn_stats. However, it looks like you forget to write the 'bn_means' and 'bn_vars' back to the model, in the current case, due to the momentum is set to 1.0, so only the mean&var of the final batch is written into the model rather than the average values of the four batches.

`
def ft_bn_stats(model, task, train_iter, device, config):

model.train()  # in order to update bn stats
model.set_bn_momentum(1.0)  # so no data leaks from one batch to the other

# compute bn stats with given task
bn_means = {n: torch.zeros_like(m.running_mean) for (n, m) in model.named_modules() if "BatchNorm" in str(type(m))}
bn_vars = {n: torch.zeros_like(m.running_var) for (n, m) in model.named_modules() if "BatchNorm" in str(type(m))}

# we don't need gradients
with torch.no_grad():
    for n, (x_t, y_t) in enumerate(train_iter):

        if n < config["EVAL"]["n_ft_bn_stats"]:
            x_t, y_t = x_t.to(device), y_t.to(device)
            model.forward(x_t, task)

            for name, module in model.named_modules():
                if "BatchNorm" in str(type(module)):
                    bn_means[name] += module.running_mean / config["EVAL"]["n_ft_bn_stats"]
                    bn_vars[name] += module.running_var / config["EVAL"]["n_ft_bn_stats"]

        else:
            break

Recommend Projects