神经网络手动实现(3)

BN层实现：

# 当前ｂａｔｃｈ均值方差计算，以及历史均值方差滑动平均计算　#
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())