在pytorch中有两个常用的模块,分别是​​nn.functional​​​和​​nn.Module​​​,这二者能够实现一些网络层的定义,对于nn中大多数的layer,在nn.functional中都有对应的​​函数算子​​。

这二者但也并非相同,还是存在区别的:

  • nn.Module中实现的layer是一个特殊的类,继承了nn.Module,例如 class Layer(nn.Module),对于继承了nn.Module的层结构,他们能够​​自动提取可学习参数​​​,并且内部已经实现好了​​forward​​函数。
  • nn.functional实现的layer是一个函数,可以直接调用,不需要实例化,对于这些函数内部是没有可学习参数的。
x = torch.randn(2, 3)

y_1 = nn.ReLU()(x)
y_2 = nn.functional.relu(x)

那么为什么同样功能要设计两个接口呢?

对于模型具有​​可学习参数​​,例如Conv2d、Linear等,最好使用nn.Module,因为继承了nn.Module能够自动提取可学习参数,也可以使用nn.functional来实现,但是这样较为复杂,需要自己手动设置参数Parameter然后传入。

如果模型​​不具备可学习参数​​,例如ReLU、MaxPool2d等,使用nn.functional和nn.Module都可以。

但是有特例,​​nn.Dropout​​,最好使用nn.Module,虽然它没有可学习参数,但是这个层有个特点就是训练和推理不同,如果使用nn.Module来实现,这时就可以使用model.train()和model.eval()来区分。

x = torch.randn(10, 4)
w = torch.randn(3, 4)
b = torch.randn(1, 3)

y = F.linear(x, w, b)
print(y.shape)

需要手动定义参数Parameter来实现,还有一种写法:

class Linear(nn.Module):
def __init__(self):
super().__init__()
self.w = torch.randn(3, 4)
self.b = torch.randn(1, 3)

def forward(self, x):
return F.linear(x, w, b)

这两种写法是一样的,对于刚入门一般都是采用下面写法,利用定义好的层,采用自定义参数写法来说能够更灵活,能够实现更为复杂的操作。

class Linear(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(4, 3)

def forward(self, x):
return self.linear(x)