Updated 5 years ago
- model.train model.eval
同时发现,如果不写这两个程序也可以运行,这是因为这两个方法是针对在网络训练和测试时采用不同方式的情况,比如Batch Normalization 和 Dropout。
训练时是正对每个min-batch的,但是在测试中往往是针对单张图片,即不存在min-batch的概念。由于网络训练完毕后参数都是固定的,因此每个批次的均值和方差都是不变的,因此直接结算所有batch的均值和方差。所有Batch Normalization的训练和测试时的操作不同
在训练中,每个隐层的神经元先乘概率P,然后在进行激活,在测试中,所有的神经元先进行激活,然后每个隐层神经元的输出乘P。
主要是针对model 在训练时和评价时不同的 Batch Normalization 和 Dropout 方法模式。
eval()时,pytorch会自动把BN和DropOut固定住,不会取平均,而是用训练好的值。
不然的话,一旦test的batch_size过小,很容易就会被BN层导致生成图片颜色失真极大。
import os
import torch
import random
import numpy as np
import torchvision
import torchvision.transforms as transforms
from torch import nn
import torch.nn.functional as F
class TorchConfig(object):
use_cuda = torch.cuda.is_available()
def __init__(self, seed=2019):
self.seed = seed
self.device = torch.device('cuda' if self.use_cuda else 'cpu') # Device configuration
self.set_seed()
# torch.set_default_tensor_type(torch.DoubleTensor)
def set_seed(self):
os.environ['PYTHONHASHSEED'] = str(self.seed)
random.seed(self.seed)
np.random.seed(self.seed)
torch.manual_seed(self.seed)
if self.use_cuda:
print('GPU: %s' % torch.cuda.get_device_name(0))
torch.cuda.manual_seed(self.seed)
torch.backends.cudnn.deterministic = True
tc = TorchConfig()
# Hyper-parameters
input_size = 784
hidden_size = 500
num_classes = 10
num_epochs = 5
batch_size = 100
learning_rate = 0.001
device = tc.device
GPU: Tesla P4
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_dataset = torchvision.datasets.MNIST(root='./data',
train=True,
transform=transforms.ToTensor())
test_dataset = torchvision.datasets.MNIST(root='./data',
train=False,
transform=transforms.ToTensor())
# Data loader (input pipeline)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=False)