import torch.nn as nn
import torch.nn.functional as F 34
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5)# submodule: Conv2d
self.conv2 = nn.Conv2d(20, 20, 5)
def forward(self, x):
x = F.relu(self.conv1(x))
return F.relu(self.conv2(x))
通过上面方式赋值的
submodule 会被注册。当调用
.cuda() 的时候,submodule 的参数也会转换为 cuda Tensor。
参数:
--
add_module(name, module)
将一个
child module
添加到当前
modle
。被添加的module可以通过name属性来获取。
例子:
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.conv = nn.Conv2d(10, 20, 4)
#和上面这个增加module的方式等价
model = Model()
print(model.conv)
输出:
Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1))children() 返回当前模型 子模块的迭代器。
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 sub_module in model.children():
print(sub_module)
输出:
Conv2d(10, 20, kernel_size=(4, 4), stride=(1, 1)) Conv2d(20, 10, kernel_size=(4, 4), stride=(1, 1))cpu(device_id=None): 将所有的模型参数 (parameters) 和 buffers 复制到 CPUNOTE :官方文档用的 move ,但我觉着 copy 更合理。 cuda(device_id=None): 将所有的模型参数 (parameters) 和 buffers 赋值 GPU 参数说明 :device_id (int, optional):如果指定的话,所有的模型参数都会复制到指定的设备上。 double(): 将 parameters 和 buffers 的数据类型转换成 double 。 eval(): 将模型设置成 evaluation 模式。 仅仅当模型中有 Dropout 和 BatchNorm是才会有影响。 float(): 将 parameters 和 buffers 的数据类型转换成 float 。 forward(* input): 定义了每次执行的 计算步骤。在所有的子类中都需要重写这个函数。 half(): 将 parameters 和 buffers 的数据类型转换成 half 。 load_state_dict(state_dict): 将 state_dict 中的 parameters 和 buffers 复制到module和它的后代中。state_dict 中的 key 必须和 model.state_dict() 返回的 key一致。NOTE:用来加载模型参数。 参数说明:state_dict (dict) –保存 parameters 和 persistent buffers 的字典。 modules(): 返回一个包含 当前模型 所有模块的迭代器。



