天天看點

model.train()&&model.eval()&&with torch.no_grad()用法

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層的行為。

注意:在測試模型時候,記得在with torch.no_grad()前面加上model.eval()   ,不加model.eval() 會導緻測試精度很低

繼續閱讀