pytorchを用いたSWA stocastic weights averageを用いた場合、BatchNormalizationはいつ,どのように組み込むのでしょうか。

NNにおける学習率について勉強しています。

stochastic weight averagingを用いるとSGDを用いた学習時にwide optimaを得られやすいとい記事を見かけました。
参考はにしているのは以下です。
https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/

SWAが何をやっているかわかりやすく以下にかかれていました。
https://towardsdatascience.com/stochastic-weight-averaging-a-new-way-to-get-state-of-the-art-results-in-deep-learning-c639ccf36a

プログラム、NNの手法についてまだまだ勉強不足なところがあり、このSWAを用いた手法についてうまく導入できないでおります。SWAの記事を見かけますと、SWAを行なった場合のweightsをそのまま用いるのは望ましくなく、BatchNormalizationを行う必要がある書かれています。
記事を引用させていただいております。

BATCH NORMALIZATION

One important detail to keep in mind is batch normalization. Batch normalization layers compute running statistics of activations during training. 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(). 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. 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.

swaを用いた際のモデルのweightsデータを保存しておいて、ネットワークにBatchNormalizationを組み込んで学習をもう一度行う必要がるのでしょうか???
どのようなタイミングでBatchNormalizationを組み込むのか、全くわからないでおります。

今調べてもわからない点は、
1.weightsデータを一旦保存して、もとのネットワークにBatchNormalizationを組み込むということなのか。
2.表を見ると,SWAを使用した時の学習率とは別の学習率を設定しているようですが、学習率も設定し直す必要があるのか
です。

以下にBatchNormalizationを用いず、SWAのコードのみを組み込んだコードを記載してみました。
kerasを使っていたのでpytorchの使い方がほぼ初めてで、コードの不十分な点があった場合すみません。

SWAを用いた際のBatchNormalizationの導入について何かアドバイスをいただけましたら幸いです。

以下テストで用いているコードです。

python
1import numpy as np
2import matplotlib.pyplot as plt
3import torch
4from torch import nn
5from torch import optim
6from torch.autograd import Variable
7
8from sklearn.datasets import load_digits
9digits=load_digits()
10X=digits.data
11y=digits.target
12
13print(X.shape) #(1797, 64)
14print(y.shape)#(1797,)
15
16from torch.utils.data import TensorDataset,DataLoader
17ds_train=TensorDataset(torch.Tensor(X_train),torch.LongTensor(y_train))
18ds_test=TensorDataset(torch.Tensor(X_test),torch.LongTensor(y_test))
19
20train_loader=DataLoader(ds_train,batch_size=100,shuffle=True)
21test_loader=DataLoader(ds_test,batch_size=100,shuffle=False)
22
23#ネットワークの設定
24from torch import nn
25 
26model = nn.Sequential()
27model.add_module('fc1', nn.Linear(64, 100))
28model.add_module('relu1', nn.ReLU())
29model.add_module('fc2', nn.Linear(100, 100))
30model.add_module('relu2', nn.ReLU())
31model.add_module('fc3', nn.Linear(100, 10))
32print(model)
33
34#クロスエントロピーを設定
35criterion = nn.CrossEntropyLoss()
36#損失関数を定義
37learning_rate=0.01
38optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
39
40
41#まずはSWAを用いない学習を行ってみる。
42def train(train_loader):
43    #訓練モードに設定
44    model.train()
45    running_loss = 0
46    for batch_idx, (images, labels) in enumerate(train_loader):
47        #view はNumPyの reshapeと似た動作。
48        #images = images.view(-1, 28 * 28)
49        
50        # 微分可能に変換  なしでも動く???
51        #images=Variable(images)
52        #labels=Variable(labels)
53        
54        # 一度計算された勾配結果を0にリセット
55        optimizer.zero_grad()
56        # 入力dataをinputし、順伝播により出力
57        outputs = model(images)
58        #lossの計算
59        loss = criterion(outputs, labels)
60        running_loss += loss.item()
61        #勾配を計算
62        loss.backward()
63        #パラメーターの更新
64        optimizer.step()
65
66    train_loss = running_loss / len(train_loader)
67    return train_loss
68
69
70def valid(test_loader):
71    
72    model.eval()
73    running_loss = 0
74    correct = 0
75    total = 0
76    #評価時は勾配は不要: with torch.no_grad()
77    with torch.no_grad():
78        for batch_idx, (images, labels) in enumerate(test_loader):
79            #images = images.view(-1, 28 * 28)
80
81            outputs = model(images)
82
83            loss = criterion(outputs, labels)
84            running_loss += loss.item()
85
86            _, predicted = torch.max(outputs, 1)
87            correct += (predicted == labels).sum().item()
88            total += labels.size(0)
89
90    val_loss = running_loss / len(test_loader)
91    val_acc = float(correct) / total
92    
93    return val_loss, val_acc
94
95
96#テストで30回epochを回してみる。
97num_epochs=30
98
99loss_list = []
100val_loss_list = []
101val_acc_list = []
102for epoch in range(num_epochs):
103    loss = train(train_loader)
104    val_loss, val_acc = valid(test_loader)
105
106    print('epoch %d, loss: %.4f val_loss: %.4f val_acc: %.4f'
107          % (epoch, loss, val_loss, val_acc))
108    
109    # logを取得
110    loss_list.append(loss)
111    val_loss_list.append(val_loss)
112    val_acc_list.append(val_acc)
113#出力：epoch 29, loss: 0.1069 val_loss: 0.1641 val_acc: 0.9537
114#学習できています。
115
116
117
118
119#以下はSWAを試しに組み込んでみました。
120
121from torchcontrib.optim import SWA
122
123criterion = nn.CrossEntropyLoss()
124learning_rate=0.01
125
126base_opt = torch.optim.SGD(model.parameters(), lr=0.1)
127#swa_startが10なので10回の学習以降からSWAを計算し始める(,,,,と自分なりに理解してます)
128optimizer = SWA(base_opt, swa_start=10, swa_freq=5, swa_lr=0.05)
129
130def train(train_loader):
131    #訓練モードに設定
132    model.train()
133    running_loss = 0
134    for batch_idx, (images, labels) in enumerate(train_loader):
135        #view はNumPyの reshapeと似た動作。
136        #images = images.view(-1, 28 * 28)
137        
138        # 微分可能に変換  なしでも動く???
139        #images=Variable(images)
140        #labels=Variable(labels)
141        
142        # 一度計算された勾配結果を0にリセット
143        optimizer.zero_grad()
144        # 入力dataをinputし、順伝播により出力
145        outputs = model(images)
146        #lossの計算
147        loss = criterion(outputs, labels)
148        running_loss += loss.item()
149        #勾配を計算
150        loss.backward()
151        #パラメーターの更新
152        optimizer.step()
153    opt.swap_swa_sgd()
154
155    train_loss = running_loss / len(train_loader)
156    return train_loss
157
158
159def valid(test_loader):
160    
161    model.eval()
162    running_loss = 0
163    correct = 0
164    total = 0
165    #評価時は勾配は不要なので with torch.no_grad()をつける。
166    with torch.no_grad():
167        for batch_idx, (images, labels) in enumerate(test_loader):
168            #images = images.view(-1, 28 * 28)
169
170            outputs = model(images)
171
172            loss = criterion(outputs, labels)
173            running_loss += loss.item()
174
175            _, predicted = torch.max(outputs, 1)
176            correct += (predicted == labels).sum().item()
177            total += labels.size(0)
178
179    val_loss = running_loss / len(test_loader)
180    val_acc = float(correct) / total
181    
182    return val_loss, val_acc
183
184
185#SWAを導入したモデルもepoch:30回にする。
186num_epochs=30
187
188loss_list = []
189val_loss_list = []
190val_acc_list = []
191for epoch in range(num_epochs):
192    loss = train(train_loader)
193    val_loss, val_acc = valid(test_loader)
194
195    print('epoch %d, loss: %.4f val_loss: %.4f val_acc: %.4f'
196          % (epoch, loss, val_loss, val_acc))
197    
198    # logging
199    loss_list.append(loss)
200    val_loss_list.append(val_loss)
201    val_acc_list.append(val_acc)
202
203#epoch 29, loss: 0.0137 val_loss: 0.1020 val_acc: 0.9611
204#何度か行なってもSWA導入前よりSWAを用いた方がaccなどスコアが良い?と言える?
205
206#この先SWAを用いた後にどうやってBatchNormalizationを用いるべきかわかりません。
207#アドバイスをいただきたいです。よろしくお願いします。
208
209コード