BN层实现:
# 当前batch均值方差计算,以及历史均值方差滑动平均计算 #
def bn(x, alpha, beta, moving_mean, moving_var, momentum, eps):
# if test #
"""
:param x: input feature tensor
:param alpha:
:param beta:
:param moving_mean:
:param moving_var: feature data variance
:param eps: 1e-5
:return:
"""
if not torch.is_grad_enabled():
# test 直接利用训练中的均值方差计算标准化输出
x_hat = (x-moving_mean)/torch.sqrt(moving_var+eps)
else:
assert len(x.shape) in (2, 4)
# if linear #
# B,F
if len(x.shape) == 2:
mean = x.mean(dim=0)
var = ((x - mean)**2).mean(dim=0)
# if conv #
# B,C,H,W
elif len(x.shape) == 4:
mean = x.mean(dim=(0, 2, 3), keepdim=True)
var = ((x - mean)**2).mean(dim=(0, 2, 3), keepdim=True)
moving_mean = momentum * moving_mean + (1-momentum) * mean
# print("mean device:", mean.device)
moving_var = momentum * moving_var + (1-momentum) * var
x_hat = (x-mean)/(torch.sqrt(var+eps))
Y = alpha * x_hat + beta
return Y, moving_mean.data, moving_var.data
# BN层构建 #
class BN(nn.Module):
def __init__(self, num_features, num_dim):
super(BN, self).__init__()
if num_dim == 2:
shape = (1, num_features)
elif num_dim == 4:
shape = (1, num_features, 1, 1)
self.alpha = nn.Parameter(torch.ones(shape))
self.beta = nn.Parameter(torch.zeros(shape))
self.moving_mean = torch.zeros(shape)
self.moving_var = torch.zeros(shape)
def forward(self, x):
# print("x device:", x.device)
if self.moving_mean.device != x.device:
self.moving_mean = self.moving_mean.to(x.device)
# print("self device:", self.moving_mean.device)
self.moving_var = self.moving_var.to(x.device)
Y, self.moving_mean, self.moving_var = bn(
x, self.alpha, self.beta,
self.moving_mean, self.moving_var, momentum=0.9, eps=1e-5)
return Y
在Lenet中加入bn层训练测试:
import torch
from d2l import torch as d2l
import torch.nn as nn
# net model #
class Reshpe(nn.Module):
def __init__(self):
super(Reshpe, self).__init__()
def forward(self, x):
return x.reshape(-1, 1, 28, 28)
class lenet(nn.Module):
def __init__(self):
super(lenet, self).__init__()
# self.a = 0
self.reshape = Reshpe()
self.net = nn.Sequential(
self.reshape, nn.Conv2d(1, 6, kernel_size=5, padding=2), BN(6, num_dim=4), nn.Sigmoid(), nn.AvgPool2d(2, 2),
nn.Conv2d(6, 16, kernel_size=5), BN(16, num_dim=4), nn.Sigmoid(), nn.AvgPool2d(2, 2),
nn.Flatten(), nn.Linear(400, 120), BN(120, num_dim=2), nn.Sigmoid(),
nn.Linear(120, 84), BN(84, num_dim=2), nn.Sigmoid(),
nn.Linear(84, 10)
)
self.init_weight()
def forward(self, x):
return self.net(x)
def init_weight(self):
for m in self.net:
if type(m) == nn.Linear or type(m) == nn.Conv2d:
nn.init.xavier_normal_(m.weight)
net = lenet()
# print(net.net)
# data #
batch_size = 128
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size=batch_size)
# loss #
loss = nn.CrossEntropyLoss()
# optim #
lr = 0.05
optim = torch.optim.SGD(net.parameters(), lr=lr)
# trian #
# """
def train_model(net, train_iter, test_iter, optim, loss, num_epochs, device):
print("train on :", device)
net.to(device)
for epoch in range(num_epochs):
metric = d2l.Accumulator(3)
net.train()
for i, (x,y) in enumerate(train_iter):
x, y = x.to(device), y.to(device)
optim.zero_grad()
y_hat = net(x)
# print("@"*15)
# print(y_hat.shape)
l = loss(y_hat, y)
l.backward()
# pdb.set_trace()
# print(net.net[1].weight.shape)
optim.step()
metric.add(l * x.shape[0], d2l.accuracy(y_hat, y), x.shape[0])
train_l = metric[0] / metric[2]
train_acc = metric[1] / metric[2]
test_acc = d2l.evaluate_accuracy_gpu(net, test_iter)
print(f"loss: {train_l:.3f}, train_acc: {train_acc:.3f}, test_acc: {test_acc:.3f}")
num_epochs = 10
train_model(net, train_iter, test_iter, optim, loss, num_epochs, d2l.try_gpu())