Python PyTorch Module用法及代碼示例

本文簡要介紹python語言中 torch.nn.Module 的用法。

用法:

class torch.nn.Module

變量：

training(bool) -布爾值表示該模塊是處於訓練模式還是評估模式。

所有神經網絡模塊的基類。

你的模型也應該繼承這個類。

模塊還可以包含其他模塊，允許將它們嵌套在樹結構中。您可以將子模塊分配為常規屬性：

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)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

以這種方式分配的子模塊將被注冊，並且當您調用 to() 等時，它們的參數也會被轉換。