how to add BatchNormalization with SWA:stochastic weights average?
I am a beginner in Deepleaning and Pytorch.

I don't understand how to use BatchNormalization in using SWA.

pytorch.org says in https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/:


  Note that the SWA averages of the weights are never used to make
  predictions during training, and so the batch normalization layers do
  not have the activation statistics computed after you reset the
  weights of your model with opt.swap_swa_sgd()


This means it's suitable for adding BatchNormalization layer after using SWA?

# it means, in my idea

#for example

opt = torchcontrib.optim.SWA(base_opt)
for i in range(100):
    opt.zero_grad()
    loss_fn(model(input), target).backward()
    opt.step()
    if i &gt; 10 and i % 5 == 0:
        opt.update_swa()
opt.swap_swa_sgd()


#save model once
torch.save(model,"swa_model.pt")

#model_load
saved_model=torch.load("swa_model.pt")

#it means adding BatchNormalization layer??
model2=saved_model
model2.add_module("Batch1",nn.BatchNorm1d(10))

# decay learning_rate more
learning_rate=0.005
optimizer = torch.optim.SGD(model2.parameters(), lr=learning_rate)

# train model again 
for epoch in range(num_epochs):
    loss = train(train_loader)
    val_loss, val_acc = valid(test_loader)


I appreciate your replying.

following your advise,

I try to make example model adding  optimizer.bn_update()

# add  optimizer.bn_update() to model

criterion = nn.CrossEntropyLoss()
learning_rate=0.01

base_opt = torch.optim.SGD(model.parameters(), lr=0.1)
optimizer = SWA(base_opt, swa_start=10, swa_freq=5, swa_lr=0.05)

def train(train_loader):
    #mode:train
    model.train()
    running_loss = 0
    for batch_idx, (images, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = model(images)
        #loss
        loss = criterion(outputs, labels)
        running_loss += loss.item()
        loss.backward()
        optimizer.step()

    optimizer.swap_swa_sgd()
    train_loss = running_loss / len(train_loader)

    return train_loss


def valid(test_loader):

    model.eval()
    running_loss = 0
    correct = 0
    total = 0
    #torch.no_grad
    with torch.no_grad():
        for batch_idx, (images, labels) in enumerate(test_loader):
            outputs = model(images)

            loss = criterion(outputs, labels)
            running_loss += loss.item()

            _, predicted = torch.max(outputs, 1)
            correct += (predicted == labels).sum().item()
            total += labels.size(0)

    val_loss = running_loss / len(test_loader)
    val_acc = float(correct) / total

    return val_loss, val_acc



num_epochs=30

loss_list = []
val_loss_list = []
val_acc_list = []
for epoch in range(num_epochs):
    loss = train(train_loader)
    val_loss, val_acc = valid(test_loader)
    optimizer.bn_update(train_loader, model)
    print('epoch %d, loss: %.4f val_loss: %.4f val_acc: %.4f'
          % (epoch, loss, val_loss, val_acc))

    # logging
    loss_list.append(loss)
    val_loss_list.append(val_loss)
    val_acc_list.append(val_acc)

# optimizer.bn_updata() 
optimizer.bn_update(train_loader, model)

# go on evaluating model,,,



What the documentation is telling you is that since SWA computes averages of weights but those weights aren't used for prediction during training the batch normalization layers won't see those weights. This means they haven't computed the respective statistics for them (as they were never able to) which is important because the weights are used during actual prediction (i.e. not during training).

This means, they assume you have batch normalization layers in your model and want to train it using SWA. This is (more or less) not straight-forward due to the reasons above.

One approach is given as follows:


  To compute the activation statistics you can just make a forward pass on your training data using the SWA model once the training is finished.


Or you can use their helper class:


  In the SWA class we provide a helper function opt.bn_update(train_loader, model). It updates the activation statistics for every batch normalization layer in the model by making a forward pass on the train_loader data loader. You only need to call this function once in the end of training.


In case you are using Pytorch's DataLoader class you can simply supply the model (after training) and the training loader to the bn_update function which updates all batch normalization statistics for you. You only need to call this function once in the end of training.



Steps to proceed:


Train your model that includes batch normalization layers using SWA
After your model has finished training, call opt.bn_update(train_loader, model) using your training data and providing your trained model

