1.model.train()
啟用 Batch Normalization 和 Dropout
如果模型中有BN層(Batch Normalization)和Dropout,需要在訓練時添加model.train()。model.train()是保證BN層能夠用到每一批資料的均值和方差。對于Dropout,model.train()是随機取一部分網絡連接配接來訓練更新參數。
用法:注意model.train()的位置
for epoch in range(num_epoches):
running_loss = torch.Tensor([0]).to(device)
pbar = tqdm(enumerate(trainloader), total=len(trainloader),position=0)
model.train()
for i, data in pbar:
path_img, img, labels = data
img, labels = img.to(device), labels.to(device)
optimizer.zero_grad()
out = model(img)
c = out.size()
loss = criterion(out, labels.long())
loss.backward()
optimizer.step()
model.train()一定要是在epoch中并且在dataloader循環前的位置上
2.model.eval
不啟用 Batch Normalization 和 Dropout
如果模型中有BN層(Batch Normalization)和Dropout,在測試時添加model.eval()。model.eval()是保證BN層能夠用全部訓練資料的均值和方差,即測試過程中要保證BN層的均值和方差不變。對于Dropout,model.eval()是利用到了所有網絡連接配接,即不進行随機舍棄神經元。
model.eval()
with torch.no_grad():
pbar = tqdm(enumerate(testloader), total=len(testloader),position=0)
for i, data in pbar:
path_img, img, labels = data
img, labels = img.to(torch.float32).to(device),
labels.to(torch.float32).to(device)
d = img.size()
e = labels.size()
out = model(img)
f = out.size()
_, pred = torch.max(out.data, 1)
3.with torch.no_grad()
with torch.no_grad()則主要是用于停止autograd子產品的工作,以起到加速和節省顯存的作用。它的作用是将該with語句包裹起來的部分停止梯度的更新,進而節省了GPU算力和顯存,但是并不會影響dropout和BN層的行為。