Neural Network模型复杂度之Batch Normalization – Python实现

import numpy
import torch
from torch import nn
from torch import optim
from torch.utils import data
from matplotlib import pyplot as plt

numpy.random.seed(0)
torch.random.manual_seed(0)


# 获取数据与封装数据
def xFunc(r, g, b):
    x = r + 2 * g + 3 * b
    return x


def yFunc(r, g, b):
    y = r ** 2 + 2 * g ** 2 + 3 * b ** 2
    return y


def lvFunc(r, g, b):
    lv = -3 * r - 4 * g - 5 * b
    return lv


class GeneDataset(data.Dataset):
    
    def __init__(self, rRange=[-1, 1], gRange=[-1, 1], bRange=[-1, 1], num=100,/
                transform=None, target_transform=None):
        self.__rRange = rRange
        self.__gRange = gRange
        self.__bRange = bRange
        self.__num = num
        self.__transform = transform
        self.__target_transform = target_transform
        
        self.__X = self.__build_X()
        self.__Y_ = self.__build_Y_()
        
        
    def __build_X(self):
        rArr = numpy.random.uniform(*self.__rRange, (self.__num, 1))
        gArr = numpy.random.uniform(*self.__gRange, (self.__num, 1))
        bArr = numpy.random.uniform(*self.__bRange, (self.__num, 1))
        X = numpy.hstack((rArr, gArr, bArr))
        return X
    
    
    def __build_Y_(self):
        rArr = self.__X[:, 0:1]
        gArr = self.__X[:, 1:2]
        bArr = self.__X[:, 2:3]
        xArr = xFunc(rArr, gArr, bArr)
        yArr = yFunc(rArr, gArr, bArr)
        lvArr = lvFunc(rArr, gArr, bArr)
        Y_ = numpy.hstack((xArr, yArr, lvArr))
        return Y_
    
    
    def __len__(self):
        return self.__num
    
    
    def __getitem__(self, idx):
        x = self.__X[idx]
        y_ = self.__Y_[idx]
        if self.__transform:
            x = self.__transform(x)
        if self.__target_transform:
            y_ = self.__target_transform(y_)
        return x, y_


# 构建模型
class Linear(nn.Module):
    
    def __init__(self, in_features, out_features, bias=True):
        super(Linear, self).__init__()
        
        self.__in_features = in_features
        self.__out_features = out_features
        self.__bias = bias
        
        self.weight = nn.Parameter(torch.randn((in_features, out_features), dtype=torch.float64))
        self.bias = nn.Parameter(torch.randn((out_features,), dtype=torch.float64))
        
        
    def forward(self, X):
        X = torch.matmul(X, self.weight)
        if self.__bias:
            X += self.bias
        return X
    
    
class Tanh(nn.Module):
    
    def __init__(self):
        super(Tanh, self).__init__()
        
        
    def forward(self, X):
        X = torch.tanh(X)
        return X
    
    
class BatchNorm(nn.Module):
    
    def __init__(self, num_features, lamda=0.9, epsilon=1.e-6):
        super(BatchNorm, self).__init__()
        
        self.__num_features = num_features
        self.__lamda = lamda
        self.__epsilon = epsilon
        self.training = True

        self.__mu_new = nn.parameter.Parameter(torch.zeros((num_features,)))
        self.__sigma_new = nn.parameter.Parameter(torch.ones((num_features,)))
        self.__mu = torch.zeros((num_features,))
        self.__sigma = torch.ones((num_features,))


    def forward(self, X):
        if self.training:
            mu_B = torch.mean(X, axis=0)
            sigma_B = torch.sqrt(torch.var(X, axis=0) + self.__epsilon)
            X = (X - mu_B) / sigma_B
            X = X * self.__sigma_new + self.__mu_new

            self.__mu = self.__lamda * self.__mu + (1 - self.__lamda) * mu_B.data
            self.__sigma = self.__lamda * self.__sigma + (1 - self.__lamda) * sigma_B.data
            return X
        else:
            X = (X - self.__mu) / self.__sigma
            X = X * self.__sigma_new + self.__mu_new
            return X


class MLP(nn.Module):

    def __init__(self, hidden_features=50, is_batch_norm=True):
        super(MLP, self).__init__()

        self.__hidden_features = hidden_features
        self.__is_batch_norm = is_batch_norm
        self.__in_features = 3
        self.__out_features = 3

        self.lin1 = Linear(self.__in_features, self.__hidden_features)
        if self.__is_batch_norm:
            self.bn1 = BatchNorm(self.__hidden_features)
        self.tanh = Tanh()
        self.lin2 = Linear(self.__hidden_features, self.__out_features)


    def forward(self, X):
        X = self.lin1(X)
        if self.__is_batch_norm:
            X = self.bn1(X)
        X = self.tanh(X)
        X = self.lin2(X)
        return X


# 构建损失函数
class MSE(nn.Module):

    def forward(self, Y, Y_):
        loss = torch.sum((Y - Y_) ** 2)
        return loss


# 训练单元与测试单元
def train_epoch(trainLoader, model, loss_fn, optimizer):
    model.train(True)

    loss = 0
    with torch.enable_grad():
        for X, Y_ in trainLoader:
            optimizer.zero_grad()

            Y = model(X)
            lossVal = loss_fn(Y, Y_)
            lossVal.backward()
            optimizer.step()

            loss += lossVal.item()
    loss /= len(trainLoader.dataset)
    return loss


def test_epoch(testLoader, model, loss_fn):
    model.train(False)

    loss = 0
    with torch.no_grad():
        for X, Y_ in testLoader:
            Y = model(X)
            lossVal = loss_fn(Y, Y_)
            loss += lossVal.item()
    loss /= len(testLoader.dataset)
    return loss


# 进行训练与测试
class BatchNormShow(object):

    def __init__(self, trainLoader, testLoader):
        self.__trainLoader = trainLoader
        self.__testLoader = testLoader


    def train(self, epochs=100):
        torch.random.manual_seed(0)
        model_BN = MLP(300, True)
        loss_BN = MSE()
        optimizer_BN = optim.Adam(model_BN.parameters(), 0.001)

        torch.random.manual_seed(0)
        model_NoBN = MLP(300, False)
        loss_NoBN = MSE()
        optimizer_NoBN = optim.Adam(model_NoBN.parameters(), 0.001)

        trainLoss_BN, testLoss_BN = self.__train_model(self.__trainLoader, self.__testLoader, /
            model_BN, loss_BN, optimizer_BN, epochs)
        trainLoss_NoBN, testLoss_NoBN = self.__train_model(self.__trainLoader, self.__testLoader, /
            model_NoBN, loss_NoBN, optimizer_NoBN, epochs)

        fig = plt.figure(figsize=(5, 4))
        ax1 = fig.add_subplot()
        ax1.plot(range(epochs), trainLoss_BN, "r-", lw=1, label="train with BN")
        ax1.plot(range(epochs), testLoss_BN, "r--", lw=1, label="test with BN")
        ax1.plot(range(epochs), trainLoss_NoBN, "b-", lw=1, label="train without BN")
        ax1.plot(range(epochs), testLoss_NoBN, "b--", lw=1, label="test without BN")
        ax1.legend()
        ax1.set(xlabel="epoch", ylabel="loss", yscale="log")
        fig.tight_layout()
        fig.savefig("batch_norm.png", dpi=100)
        plt.show()


    def __train_model(self, trainLoader, testLoader, model, loss_fn, optimizer, epochs):
        trainLossList = list()
        testLossList = list()

        for epoch in range(epochs):
            trainLoss = train_epoch(trainLoader, model, loss_fn, optimizer)
            testLoss = test_epoch(testLoader, model, loss_fn)
            trainLossList.append(trainLoss)
            testLossList.append(testLoss)
            print(epoch, trainLoss, testLoss)
        return trainLossList, testLossList



if __name__ == "__main__":
    trainData = GeneDataset([-1, 1], [-1, 1], [-1, 1], num=1000, /
        transform=torch.tensor, target_transform=torch.tensor)
    testData = GeneDataset([-1, 1], [-1, 1], [-1, 1], num=300, /
        transform=torch.tensor, target_transform=torch.tensor)
    trainLoader = data.DataLoader(trainData, batch_size=len(trainData), shuffle=False)
    testLoader = data.DataLoader(testData, batch_size=len(testData), shuffle=False)
    bnsObj = BatchNormShow(trainLoader, testLoader)
    epochs = 10000
    bnsObj.train(epochs)