nn.Module的深入分析

nn.Module作为基类,其结构如下所示:

def __init__(self):
self._parameters = OrderedDict()   ##用来保存用户设置的参数字典,key:params Value:用户自定义的关键字和值。系统定义的不存与此。
self._modules=OrderedDict()  ##存放子module,使用submodule生成的module会存于此。
self._buffers=OrderedDict() ##存放一些缓存结果。
self._backward=OrderedDict()  
self._forward=OrderedDict()   
self.training=True  ##dropout在训练和测试中采取的模式不同,通过training决定前向传播策略。

调用model.train()会把所有的module设置为训练模式,调用model.eval()会把所有的training属性设置为False,

钩子函数hook可以对得到对某一层的输出。格式如下:

model=VGG();
features=t.Tensor();  ##生成一个tensor用于存储后面得到的数据
def hook(module,inputs,output)
  ##把这层的输出复制到features中
  features.copy_(output.data)
handle=model.layer8.register_forward_hook(hook)  ##把新定义的hook注册进去
__=model(input)  ##得到的结果在features里面
##hook使用完毕后删除,否则后面调用model(input)时调用到layer8就不向后输出了
handle.remove()

保存模型的两种方式:

import torch as t
t.save(net.state_dict(),'net.pth');
model.load_state_dic(t.load('net.pth'))

#虽然第二种要依赖路径,但是个人感觉第二种更简便,嘻嘻~
t.save(net,'./net.pth');
net2=net();
net2.load('net.pth');

将模型和数据放在GPU上运行,可以直接使用model=mode.cuda(),input=input.cuda().将数据转存在gpu上运行。

pytorch提供了两个函数,可以使在多个GPU上运行。

nn.parallel.data_parallel(module,inputs,device_ids=None,output_device=None,dim=0,module_kwarge=None)
##或者是
##直接返回一个新的module,自动在多GPU上进行并行加速
class torch.nn.DataParallel(module,device_ids=None,output_device=None,dim=0)

通过设置device_ids来确定在哪个GPU上进行优化,设置output_device来确定把结果输出在哪一块,一般GPU都会分为好几块,跟在运行的时候前面加参数CUDA_VISIBLE_DEVICE=1,在GPU编号为1的块上运行效果一样,只是上面的更加方便,可以使用[0,1]这样确定多个使用的GPU编号。例如:

###DataParallel是直接将一个batch的数据进行均分送到各个GPU上,再把GPU的结果收集。
new_net=nn.DataParallel(net,device_ids=[0,1]);
output=new_new(inputs);

##或者是
###自动通过多并行计算得到结果
output=nn.parallel.data_parallel(net,inputs,device_ids=[0,1])

你可能感兴趣的:(nn.Module的深入分析)