import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.add_module("conv", nn.Conv2d(10, 20, 4))
self.add_module("conv1", nn.Conv2d(20 ,10, 4))
model = Model()
for module in model.modules():
print(module)
Model (
(conv): Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))
(conv1): Conv2d(20, 10, kernel_size=(4, 4), stride=(1, 1))
)
Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))
Conv2d(20, 10, kernel_size=(4, 4), stride=(1, 1))
可以看出,modules()返回的iterator不止包含子模块。这是和childern()的不同。NOTE:重复的模块只被返回一次(children()也是)。在下面的例子中submodule只会被返回一次。
import torch.nn as nn
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
submodule = nn.Conv2d(10, 20, 4)
self.add_module("conv", submodule)
self.add_module("conv1", submodule)
model = Model()
for module in model.modules():
print(module)
Model (
(conv): Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1)) , →
(conv1): Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1)) , →
)
Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))
named_children()
返回包含模型当前子模块的迭代器,yield模块名字和模块本身。 例子:
for name, module in model.named_children():
if name in ['conv4', 'conv5']:
print(module)
for param in model.parameters():
print(type(param.data), param.size())
(20L,)
(20L, 1L, 5L, 5L)
register_backward_hook(hook)
self.register_buffer('running_mean',
torch.zeros(num_features))
self.running_mean