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

前言:

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

self.modules()

首先来讲一下nn.Module类中的一个方法:self.modules(), 他会返回该网络中的所有modules.

具体代码示例如下

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_()