文章来源:https://www.toymoban.com/news/detail-726683.html
文章来源地址https://www.toymoban.com/news/detail-726683.html
from pathlib import Path
from typing import Optional
import torch
import torch.nn as nn
from torch import Tensor
class BN(nn.Module):
def __init__(self,num_features,momentum=0.1,eps=1e-8):##num_features是通道数
"""
初始化方法
:param num_features:特征属性的数量,也就是通道数目C
"""
super(BN, self).__init__()
##register_buffer:将属性当成parameter进行处理,唯一的区别就是不参与反向传播的梯度求解
self.register_buffer('running_mean', torch.zeros(1, num_features, 1, 1))
self.register_buffer('running_var', torch.zeros(1, num_features, 1, 1))
self.running_mean: Optional[Tensor]
self.running_var: Optional[Tensor]
self.running_mean=torch.zeros([1,num_features,1,1])
self.running_var=torch.zeros([1,num_features,1,1])
self.gamma=nn.Parameter(torch.ones([1,num_features,1,1]))
self.beta=nn.Parameter(torch.zeros(1,num_features,1,1))
self.eps=eps
self.momentum=momentum
def forward(self,x):
"""
前向过程
output=(x-μ)/α*γ+β
:param x: [N,C,H,W]
:return: [N,C,H,W]
"""
if self.training:
#训练阶段--》使用当前批次的数据
_mean=torch.mean(x,dim=(0,2,3),keepdim=True)
_var = torch.var(x, dim=(0, 2, 3), keepdim=True)
#将训练过程中的均值和方差保存下来--方便推理的时候使用--》滑动平均
self.running_mean=self.momentum*self.running_mean+(1.0-self.momentum)*_mean
self.running_var=self.momentum*self.running_var+(1.0-self.momentum)*_var
else:
#推理阶段-->使用的是训练过程中的累积数据
_mean=self.running_mean
_var=self.running_var
z=(x-_mean)/torch.sqrt(_var+self.eps)*self.gamma+self.beta
return z
if __name__ == '__main__':
torch.manual_seed(28)
path_dir=Path("./output/models")
path_dir.mkdir(parents=True,exist_ok=True)
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
bn=BN(num_features=12)
bn.to(device)#只针对子模块和参数进行转换
#模拟训练过程
bn.train()
xs=[torch.randn(8,12,32,32).to(device) for _ in range(10)]
for _x in xs:
bn(_x)
print(bn.running_mean.view(-1))
print(bn.running_var.view(-1))
#模拟推理过程
bn.eval()
_r=bn(xs[0])
print(_r.shape)
bn=bn.cpu()#保存都是以cpu保存,恢复再自己转回GPU上
#模拟模型保存
torch.save(bn,str(path_dir/'bn_model.pkl'))
#state_dict:获取当前模块的所有参数(Parameter+register_buffer)
torch.save(bn.state_dict(),str(path_dir/"bn_params.pkl"))
#pt结构的保存
traced_script_module=torch.jit.trace(bn.eval(),xs[0].cpu())
traced_script_module.save("./output/bn_model.pt")
#模拟模型恢复
bn_model=torch.load(str(path_dir/"bn_model.pkl"),map_location='cpu')
bn_params=torch.load(str(path_dir/"bn_params.pkl"),map_location='cpu')
print(len(bn_params))
到了这里,关于BN体系理解——类封装复现的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!