PyTorch中self.layers的使用小結

更新時間：2024年01月26日 09:53:51 作者：風箏超冷

self.layers?是一個用于存儲網絡層的屬性,本文主要介紹了PyTorch中self.layers的使用小結,文中通過示例代碼介紹的非常詳細,對大家的學習或者工作具有一定的參考學習價值,需要的朋友們下面隨著小編來一起學習學習吧

為什么使用 nn.ModuleList？

在PyTorch中，當需要處理多個神經網絡層時，通常使用 nn.ModuleList 或 nn.Sequential。這些容器類能夠確保其中包含的所有模塊（層）都被正確注冊，這樣PyTorch就可以跟蹤它們的參數(shù)，實現(xiàn)自動梯度計算和參數(shù)更新。

self.layers的作用

class UserDefined(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
            ]))
    
    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x
        return x

在自定義的類中，self.layers 具有以下特點和作用：

存儲層: 它存儲了Transformer模型中所有的層。在這個例子中，每層由一個預歸一化的多頭注意力模塊和一個預歸一化的前饋網絡模塊組成。
動態(tài)創(chuàng)建層: 通過在 for 循環(huán)中添加層，self.layers 能夠根據(jù)提供的 depth 參數(shù)動態(tài)創(chuàng)建相應數(shù)量的Transformer層。
維護層順序: nn.ModuleList 維護了添加到其中的模塊的順序，這對于保持層的順序非常重要，因為在Transformer模型中數(shù)據(jù)需要按照特定的順序通過這些層。
模型前向傳播: 在 forward 方法中，self.layers 被遍歷，數(shù)據(jù)依次通過每一層。這個過程涉及到每層中多頭注意力和前饋網絡的計算。