torch 变量_torch源码阅读之module.py

这一周将必看的torch文件全部看完！！！！

在构建神经网络架构时，都要继承自基类nn.Module。

class Module(object)    
 def __init__(self):
        self._backend = thnn_backend  
        self._parameters = OrderedDict()   # 字典，保存用户直接设置的parameters
        self._buffers = OrderedDict()      # 缓存
        self._backward_hooks = OrderedDict()  # 钩子技术，用来提取中间变量
        self._forward_hooks = OrderedDict()
        self._forward_pre_hooks = OrderedDict()
        self._modules = OrderedDict()
        self.training = True               # BatchNorm和DropOut层在训练阶段和测试阶段中采取的策略不同，通过判断traing值决定前向传播的策略

接下来看看该类的比较常用的几个function：

def forward(self, *input):
       # 继承该类的所有子类中都要重写该函数
    raise NotImplementedError

增加子模块到当前模块中：

    def add_module(self, name, module):

        if not isinstance(module, Module) and module is not None:
            raise TypeError("{} is not a Module subclass".format(
                torch.typename(module)))
        elif hasattr(self, name) and name not in self._modules:
            raise KeyError("attribute '{}' already exists".format(name))
        elif '.' in name:
            raise KeyError("module name can't contain "."")
        elif name == '':
            raise KeyError("module name can't be empty string """)
        self._modules[name] = module

该函数的最后一行self._modules[name] = module，而self._module是一个字典，保存的是子module。函数children可以查看直接子module。函数modules可以查看所有的子module。

    def children(self):
        for name, module in self.named_children():
            yield module       

    def named_children(self):
        memo = set()
        for name, module in self._modules.items():
            if module is not None and module not in memo:
                memo.add(module)
                yield name, module   # yield同return差不多，但是它的type是一个生成器，可以节省时间

函数modules可以查看所有的子module：

    def modules(self):
        for name, module in self.named_modules():
            yield module

    def named_modules(self, memo=None, prefix=''):
 
        if memo is None:
            memo = set()
        if self not in memo:
            memo.add(self)
            yield prefix, self
            for name, module in self._modules.items():
                if module is None:
                    continue
                submodule_prefix = prefix + ('.' if prefix else '') + name
                for m in module.named_modules(memo, submodule_prefix):
                    yield m

对于batchnorm, dropout, instancenorm等在训练和测试阶段行为差距较大的层，如果在测试时不将其training值设为False，则会有很大影响。推荐调用model.train()函数，它会将当前module及子module中的所有training属性都设为True。相应的，model.eval()函数会把training属性都设为False。

    def train(self, mode=True):
        self.training = mode
        for module in self.children():
            module.train(mode)
        return self

     def eval(self):
        return self.train(False)

可以用named_parameters()来查询子module的名字和params的大小等信息：

    def parameters(self):
        for name, param in self.named_parameters():
            yield param

    def named_parameters(self, memo=None, prefix=''):
        
        if memo is None:
            memo = set()
        for name, p in self._parameters.items():
            if p is not None and p not in memo:
                memo.add(p)
                yield prefix + ('.' if prefix else '') + name, p
        for mname, module in self.named_children():
            submodule_prefix = prefix + ('.' if prefix else '') + mname
            for name, p in module.named_parameters(memo, submodule_prefix):
                yield name, p

使用方法：

for name, param in net.named_parameters():
    print(name, param.size())

将模型放到GPU上运行也十分的简单，调用cuda函数就好，下面是在单GPU上运行：

    def cuda(self, device=None):
        return self._apply(lambda t: t.cuda(device))

model = model.cuda()  # 将模型的所有参数转存到GPU上
input.cuda()   # 将输入数据放置在GPU上

多GPU上并行计算时，pytorch提供了其他两个函数：nn.parallel.data_parallel()和class torch.nn.DataParallel()，这个以后再说。

加更：2020.04.10

self.load_state_dict()加载指定路径的模型。

    def load_state_dict(self, state_dict, strict=True):
        missing_keys = []
        unexpected_keys = []
        error_msgs = []

        # copy state_dict so _load_from_state_dict can modify it
        metadata = getattr(state_dict, '_metadata', None)
        state_dict = state_dict.copy()
        if metadata is not None:
            state_dict._metadata = metadata

        def load(module, prefix=''):
            module._load_from_state_dict(
                state_dict, prefix, strict, missing_keys, unexpected_keys, error_msgs)
            for name, child in module._modules.items():
                if child is not None:
                    load(child, prefix + name + '.')

        load(self)

        if strict:
            error_msg = ''
            if len(unexpected_keys) > 0:
                error_msgs.insert(
                    0, 'Unexpected key(s) in state_dict: {}. '.format(
                        ', '.join('"{}"'.format(k) for k in unexpected_keys)))
            if len(missing_keys) > 0:
                error_msgs.insert(
                    0, 'Missing key(s) in state_dict: {}. '.format(
                        ', '.join('"{}"'.format(k) for k in missing_keys)))

        if len(error_msgs) > 0:
            raise RuntimeError('Error(s) in loading state_dict for {}:nt{}'.format(
                               self.__class__.__name__, "nt".join(error_msgs)))