pytorch中利用self.module()方法来初始化模型权重
程序员文章站
2022-07-14 20:17:43
...
前言:
每次在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_()