欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

pytorch中利用self.module()方法来初始化模型权重

程序员文章站 2022-07-14 20:17:43
...

前言:

每次在pytorch中构建完模型之后都要在另外初始化方法, 非常的麻烦, 因此直接利用构建模型类中的初始化方法 __ init __ ()以及继承的方法self.modules()来初始化模型权重.

self.modules()

首先来讲一下nn.Module类中的一个方法:self.modules(), 他会返回该网络中的所有modules.
pytorch中利用self.module()方法来初始化模型权重

具体代码示例如下

class Network(nn.Module):
	def __init__(self):
		supe().__init__()
		self.Conv2d = nn.Conv2d(3, 10)
		sefl.bn = nn.BatchNorm2d(10)
		self.relu = nn.ReLU()
		self._init_weight()  #在初始化网络时, 会执行该函数,然后初始化网络中的每个module
		
	def forward(self, x):
		x = self.Conv2d(x)
		x = self.bn(x)
		return self.relu(x)
		
	def _init_weight(self):
		for m in self.modules()   #继承nn.Module的方法
			if isinstance(m, nn.Conv2d):
				torch.nn.init.kaiming_normal_(m.weight)
			elif isinstance(m, nn.BatchNorm2d):
				m.weight.data.fill_(1)
				m.bias.data.zero_()