分类目录:《深入浅出Pytorch函数》总目录
torch.nn.Module
是所有Pytorch中所有神经网络模型的基类,我们的神经网络模型也应该继承这个类。Modules
可以包含其它Modules
,也允许使用树结构嵌入他们,还可以将子模块赋值给模型属性。
实例
import torch.nn as nn
import torch.nn.functional as F
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
。
函数
eval()
将模块设置为evaluation
模式,相当于self.train(False)
。这个函数仅当模型中有Dropout
和BatchNorm
时才会有影响。
def eval(self: T) -> T:
r"""Sets the module in evaluation mode.
This has any effect only on certain modules. See documentations of
particular modules for details of their behaviors in training/evaluation
mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,
etc.
This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.
See :ref:`locally-disable-grad-doc` for a comparison between
`.eval()` and several similar mechanisms that may be confused with it.
Returns:
Module: self
"""
return self.train(False)