以下代码是对conv+bn的融合，这里，均值、方差都是对batchsize,height,width这三个维度求的，因此 m e a n , v a r , γ , β 都 是 一 个 N 维 向 量 ， 其 中 N 为 卷 积 层 输 出 特 征 的 通 道 数 mean,var,\gamma,\beta都是一个N维向量，其中N为卷积层输出特征的通道数 mean,var,γ,β都是一个N维向量，其中N为卷积层输出特征的通道数
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
def myconv(input,weight,bias): #S=1,Pad=0
B,N,Hin,Win=input.shape
M,N,Kx,Ky=weight.shape
Hout=(Hin-Kx+1)
Wout=(Win-Ky+1)
output=np.zeros((B,M,Hout,Wout))
for i in range(Hout):
for j in range(Wout):
for k in range(M):
for l in range(B):
output[l,k,i,j]=np.sum(weight[k,:,:,:]*input[l,:,i:i+Kx,j:j+Ky])+bias[k]
return output
def mypool(input): #S=2,K=2
B,N,Hin,Win=input.shape
Hout=Hin//2
Wout=Win//2
output=np.zeros((B,N,Hout,Wout))
for i in range(Hout):
for j in range(Wout):
for k in range(N):
for l in range(B):
output[l,k,i,j]=np.max(input[l,k,2*i:2*i+2,2*j:2*j+2])
return output
def mybn(input,bn_mean,bn_var,bn_gama,bn_beta):
B,N,H,W=input.shape
#(x-mean)/sqrt(var+eps)*gama+beta
bn_mean=bn_mean.reshape(1,-1,1,1)
bn_var=bn_var.reshape(1,-1,1,1)
bn_gama=bn_gama.reshape(1,-1,1,1)
bn_beta=bn_beta.reshape(1,-1,1,1)
output=np.zeros_like(input)
output=(input-bn_mean)/np.sqrt(bn_var)*bn_gama+bn_beta
return output
def myrelu(input):
return np.maximum(0,input)
def myconv_bn(input,weight,bias,bn_mean,bn_var,bn_gama,bn_beta): #bias(M,),bn_mean(M,),bn_var(M,),bn_gama(M,),bn_beta(M,)
alpha=bn_gama/np.sqrt(bn_var)
weight=weight*alpha.reshape(-1,1,1,1)
bias=(bias-bn_mean)*alpha+bn_beta
return myconv(input,weight,bias)
#*****************************************************************************************************************************************************************
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4 * 4 * 50, 500)
self.fc2 = nn.Linear(500, 10)
self.bn1 = nn.BatchNorm2d(20, momentum=0.01)
self.bn2 = nn.BatchNorm2d(50, momentum=0.01)
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.bn2(self.conv2(x)))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4 * 4 * 50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.log_softmax(x, dim=1)
def train(model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
def test( model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
torch.manual_seed(1)
device = torch.device("cuda")
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=128, shuffle=True)
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.005, momentum=0.5)
for epoch in range(1, 2 + 1):
train(model, device, train_loader, optimizer, epoch)
test(model, device, test_loader)
Wc1=model.state_dict()['conv1.weight'].cpu().numpy()
bc1=model.state_dict()['conv1.bias'].cpu().numpy()
Wc2=model.state_dict()['conv2.weight'].cpu().numpy()
bc2=model.state_dict()['conv2.bias'].cpu().numpy()
Wf1=model.state_dict()['fc1.weight'].cpu().numpy()
bf1=model.state_dict()['fc1.bias'].cpu().numpy()
Wf2=model.state_dict()['fc2.weight'].cpu().numpy()
bf2=model.state_dict()['fc2.bias'].cpu().numpy()
gama1=model.state_dict()['bn1.weight'].cpu().numpy()
beta1=model.state_dict()['bn1.bias'].cpu().numpy()
mean1=model.state_dict()['bn1.running_mean'].cpu().numpy()
var1=model.state_dict()['bn1.running_var'].cpu().numpy()
gama2=model.state_dict()['bn2.weight'].cpu().numpy()
beta2=model.state_dict()['bn2.bias'].cpu().numpy()
mean2=model.state_dict()['bn2.running_mean'].cpu().numpy()
var2=model.state_dict()['bn2.running_var'].cpu().numpy()
print(Wc1.shape)
print(bc1.shape)
print(Wc2.shape)
print(bc2.shape)
print(Wf1.shape)
print(bf1.shape)
print(Wf2.shape)
print(bf2.shape)
print(gama1.shape)
print(beta1.shape)
print(mean1.shape)
print(var1.shape)
print(gama2.shape)
print(beta2.shape)
print(mean2.shape)
print(var2.shape)
for batch_idx, (data, target) in enumerate(train_loader):
input=data.numpy()
x=myconv_bn(input,Wc1,bc1,mean1,var1,gama1,beta1)
x=myrelu(x)
x=mypool(x)
x=myconv_bn(x,Wc2,bc2,mean2,var2,gama2,beta2)
x=myrelu(x)
x=mypool(x)
x=x.reshape(128,-1)
x=np.dot(x,Wf1.T)+np.expand_dims(bf1,0).repeat(128,axis=0)
x=myrelu(x)
x=np.dot(x,Wf2.T)+np.expand_dims(bf2,0).repeat(128,axis=0)
output=np.argmax(x,axis=1)
label=target.numpy()
correct=0
for i in range(128):
if label[i]==output[i]:
correct+=1
print(correct/128)
