NeuralNetwork模型复杂度之BatchNormalizationPython实现

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)