Pytorch how and when to use Module Sequential ModuleList and ModuleDict
1.0.0
Pytorch 1.5 更新
您可以在這裡找到程式碼
Pytorch 是一個開源深度學習框架，提供建立 ML 模型的智慧方法。即使文件製作精良，我仍然發現大多數人都沒有在 PyTorch 中編寫良好且有組織的程式碼。
今天，我們將了解如何使用 PyTorch 的三個主要構建塊： Module, Sequential and ModuleList 。我們將從一個範例開始，然後迭代地使其變得更好。
所有這四個類別都包含在torch.nn中
import torch . nn as nn
# nn.Module
# nn.Sequential
# nn.Module 模組是主要構建塊，它定義了所有神經網路的基類，您必須對其進行子類化。
讓我們建立一個經典的 CNN 分類器作為範例：
import torch . nn . functional as F
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , n_classes ):
super (). __init__ ()
self . conv1 = nn . Conv2d ( in_c , 32 , kernel_size = 3 , stride = 1 , padding = 1 )
self . bn1 = nn . BatchNorm2d ( 32 )
self . conv2 = nn . Conv2d ( 32 , 64 , kernel_size = 3 , stride = 1 , padding = 1 )
self . bn2 = nn . BatchNorm2d ( 64 )
self . fc1 = nn . Linear ( 64 * 28 * 28 , 1024 )
self . fc2 = nn . Linear ( 1024 , n_classes )
def forward ( self , x ):
x = self . conv1 ( x )
x = self . bn1 ( x )
x = F . relu ( x )
x = self . conv2 ( x )
x = self . bn2 ( x )
x = F . relu ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . fc1 ( x )
x = F . sigmoid ( x )
x = self . fc2 ( x )
return x model = MyCNNClassifier ( 1 , 10 )
print ( model ) MyCNNClassifier(
(conv1): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(bn1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(fc1): Linear(in_features=50176, out_features=1024, bias=True)
(fc2): Linear(in_features=1024, out_features=10, bias=True)
)
這是一個非常簡單的分類器，其編碼部分使用具有 3x3 卷積 + batchnorm + relu 的兩個層，以及具有兩個線性層的解碼部分。如果您對 PyTorch 不陌生，您可能以前見過這種類型的編碼，但有兩個問題。
如果我們想要新增一層，我們必須再次在__init__和forward函數中編寫大量程式碼。另外，如果我們有一些公共模組想要在另一個模型中使用，例如 3x3 conv + batchnorm + relu，我們必須重新編寫它。
Sequential 是一個可以堆疊在一起並同時運作的模組的容器。
您可以注意到我們必須將所有內容儲存到self中。我們可以使用Sequential來改進我們的程式碼。
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , n_classes ):
super (). __init__ ()
self . conv_block1 = nn . Sequential (
nn . Conv2d ( in_c , 32 , kernel_size = 3 , stride = 1 , padding = 1 ),
nn . BatchNorm2d ( 32 ),
nn . ReLU ()
)
self . conv_block2 = nn . Sequential (
nn . Conv2d ( 32 , 64 , kernel_size = 3 , stride = 1 , padding = 1 ),
nn . BatchNorm2d ( 64 ),
nn . ReLU ()
)
self . decoder = nn . Sequential (
nn . Linear ( 64 * 28 * 28 , 1024 ),
nn . Sigmoid (),
nn . Linear ( 1024 , n_classes )
)
def forward ( self , x ):
x = self . conv_block1 ( x )
x = self . conv_block2 ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , 10 )
print ( model ) MyCNNClassifier(
(conv_block1): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(conv_block2): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(decoder): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
(2): Linear(in_features=1024, out_features=10, bias=True)
)
)
好多了呃？
您是否注意到conv_block1和conv_block2看起來幾乎一樣？我們可以建立一個傳回nn.Sequential的函數來簡化程式碼！
def conv_block ( in_f , out_f , * args , ** kwargs ):
return nn . Sequential (
nn . Conv2d ( in_f , out_f , * args , ** kwargs ),
nn . BatchNorm2d ( out_f ),
nn . ReLU ()
)然後我們就可以在我們的模組中呼叫這個函數
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , n_classes ):
super (). __init__ ()
self . conv_block1 = conv_block ( in_c , 32 , kernel_size = 3 , padding = 1 )
self . conv_block2 = conv_block ( 32 , 64 , kernel_size = 3 , padding = 1 )
self . decoder = nn . Sequential (
nn . Linear ( 64 * 28 * 28 , 1024 ),
nn . Sigmoid (),
nn . Linear ( 1024 , n_classes )
)
def forward ( self , x ):
x = self . conv_block1 ( x )
x = self . conv_block2 ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , 10 )
print ( model ) MyCNNClassifier(
(conv_block1): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(conv_block2): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(decoder): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
(2): Linear(in_features=1024, out_features=10, bias=True)
)
)
更乾淨！仍然conv_block1和conv_block2幾乎相同！我們可以使用nn.Sequential合併它們
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , n_classes ):
super (). __init__ ()
self . encoder = nn . Sequential (
conv_block ( in_c , 32 , kernel_size = 3 , padding = 1 ),
conv_block ( 32 , 64 , kernel_size = 3 , padding = 1 )
)
self . decoder = nn . Sequential (
nn . Linear ( 64 * 28 * 28 , 1024 ),
nn . Sigmoid (),
nn . Linear ( 1024 , n_classes )
)
def forward ( self , x ):
x = self . encoder ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , 10 )
print ( model ) MyCNNClassifier(
(encoder): Sequential(
(0): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(1): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
(decoder): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
(2): Linear(in_features=1024, out_features=10, bias=True)
)
)
self.encoder現在持有 Booth conv_block 。我們為模型解耦了邏輯，使其更易於閱讀和重複使用。我們的conv_block函數可以導入並在另一個模型中使用。
如果我們可以在self.encoder中新增圖層怎麼辦，對它們進行硬編碼並不方便：
self . encoder = nn . Sequential (
conv_block ( in_c , 32 , kernel_size = 3 , padding = 1 ),
conv_block ( 32 , 64 , kernel_size = 3 , padding = 1 ),
conv_block ( 64 , 128 , kernel_size = 3 , padding = 1 ),
conv_block ( 128 , 256 , kernel_size = 3 , padding = 1 ),
)如果我們可以將大小定義為一個陣列並自動建立所有層而不需要編寫每一層，那該有多好？幸運的是，我們可以建立一個陣列並將其傳遞給Sequential
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , n_classes ):
super (). __init__ ()
self . enc_sizes = [ in_c , 32 , 64 ]
conv_blocks = [ conv_block ( in_f , out_f , kernel_size = 3 , padding = 1 )
for in_f , out_f in zip ( self . enc_sizes , self . enc_sizes [ 1 :])]
self . encoder = nn . Sequential ( * conv_blocks )
self . decoder = nn . Sequential (
nn . Linear ( 64 * 28 * 28 , 1024 ),
nn . Sigmoid (),
nn . Linear ( 1024 , n_classes )
)
def forward ( self , x ):
x = self . encoder ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , 10 )
print ( model ) MyCNNClassifier(
(encoder): Sequential(
(0): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(1): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
(decoder): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
(2): Linear(in_features=1024, out_features=10, bias=True)
)
)
讓我們來分解一下。我們建立了一個陣列self.enc_sizes來保存編碼器的大小。然後我們透過迭代大小來建立一個陣列conv_blocks 。由於我們必須為每一層指定展位的 in 尺寸和 outsize，因此我們透過將 size' 陣列移動一位來將其自身zip 。
為了清楚起見，請看以下範例：
sizes = [ 1 , 32 , 64 ]
for in_f , out_f in zip ( sizes , sizes [ 1 :]):
print ( in_f , out_f ) 1 32
32 64
然後，由於Sequential不接受列表，因此我們使用*運算子對其進行分解。
田田！現在，如果我們只想添加尺寸，我們可以輕鬆地將新數字添加到清單中。將尺寸作為參數是常見的做法。
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , enc_sizes , n_classes ):
super (). __init__ ()
self . enc_sizes = [ in_c , * enc_sizes ]
conv_blocks = [ conv_block ( in_f , out_f , kernel_size = 3 , padding = 1 )
for in_f , out_f in zip ( self . enc_sizes , self . enc_sizes [ 1 :])]
self . encoder = nn . Sequential ( * conv_blocks )
self . decoder = nn . Sequential (
nn . Linear ( 64 * 28 * 28 , 1024 ),
nn . Sigmoid (),
nn . Linear ( 1024 , n_classes )
)
def forward ( self , x ):
x = self . encoder ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , [ 32 , 64 , 128 ], 10 )
print ( model ) MyCNNClassifier(
(encoder): Sequential(
(0): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(1): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(2): Sequential(
(0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
(decoder): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
(2): Linear(in_features=1024, out_features=10, bias=True)
)
)
我們可以對解碼器部分做同樣的事情
def dec_block ( in_f , out_f ):
return nn . Sequential (
nn . Linear ( in_f , out_f ),
nn . Sigmoid ()
)
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , enc_sizes , dec_sizes , n_classes ):
super (). __init__ ()
self . enc_sizes = [ in_c , * enc_sizes ]
self . dec_sizes = [ 64 * 28 * 28 , * dec_sizes ]
conv_blocks = [ conv_block ( in_f , out_f , kernel_size = 3 , padding = 1 )
for in_f , out_f in zip ( self . enc_sizes , self . enc_sizes [ 1 :])]
self . encoder = nn . Sequential ( * conv_blocks )
dec_blocks = [ dec_block ( in_f , out_f )
for in_f , out_f in zip ( self . dec_sizes , self . dec_sizes [ 1 :])]
self . decoder = nn . Sequential ( * dec_blocks )
self . last = nn . Linear ( self . dec_sizes [ - 1 ], n_classes )
def forward ( self , x ):
x = self . encoder ( x )
x = x . view ( x . size ( 0 ), - 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , [ 32 , 64 ], [ 1024 , 512 ], 10 )
print ( model ) MyCNNClassifier(
(encoder): Sequential(
(0): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(1): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
(decoder): Sequential(
(0): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
)
(1): Sequential(
(0): Linear(in_features=1024, out_features=512, bias=True)
(1): Sigmoid()
)
)
(last): Linear(in_features=512, out_features=10, bias=True)
)
我們遵循相同的模式，為解碼部分建立一個新區塊，線性 + sigmoid，並傳遞一個具有大小的陣列。我們必須添加self.last因為我們不想啟動輸出
現在，我們甚至可以將模型一分為二！編碼器+解碼器
class MyEncoder ( nn . Module ):
def __init__ ( self , enc_sizes ):
super (). __init__ ()
self . conv_blocks = nn . Sequential ( * [ conv_block ( in_f , out_f , kernel_size = 3 , padding = 1 )
for in_f , out_f in zip ( enc_sizes , enc_sizes [ 1 :])])
def forward ( self , x ):
return self . conv_blocks ( x )
class MyDecoder ( nn . Module ):
def __init__ ( self , dec_sizes , n_classes ):
super (). __init__ ()
self . dec_blocks = nn . Sequential ( * [ dec_block ( in_f , out_f )
for in_f , out_f in zip ( dec_sizes , dec_sizes [ 1 :])])
self . last = nn . Linear ( dec_sizes [ - 1 ], n_classes )
def forward ( self , x ):
return self . dec_blocks ()
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , enc_sizes , dec_sizes , n_classes ):
super (). __init__ ()
self . enc_sizes = [ in_c , * enc_sizes ]
self . dec_sizes = [ self . enc_sizes [ - 1 ] * 28 * 28 , * dec_sizes ]
self . encoder = MyEncoder ( self . enc_sizes )
self . decoder = MyDecoder ( self . dec_sizes , n_classes )
def forward ( self , x ):
x = self . encoder ( x )
x = x . flatten ( 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , [ 32 , 64 ], [ 1024 , 512 ], 10 )
print ( model ) MyCNNClassifier(
(encoder): MyEncoder(
(conv_blocks): Sequential(
(0): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
(1): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
)
)
(decoder): MyDecoder(
(dec_blocks): Sequential(
(0): Sequential(
(0): Linear(in_features=50176, out_features=1024, bias=True)
(1): Sigmoid()
)
(1): Sequential(
(0): Linear(in_features=1024, out_features=512, bias=True)
(1): Sigmoid()
)
)
(last): Linear(in_features=512, out_features=10, bias=True)
)
)
請注意， MyEncoder和MyDecoder也可以是傳回nn.Sequential的函數。我更喜歡將第一個模式用於模型，將第二個模式用於建立區塊。
透過將我們的模組分成子模組，可以更輕鬆地共享程式碼、調試和測試它。
ModuleList允許您將Module儲存為列表。當您需要遍歷圖層並儲存/使用某些資訊（例如在 U-net 中）時，它會很有用。
Sequential的主要差異在於ModuleList沒有forward方法，因此內層沒有連接。假設我們需要解碼器中每一層的每個輸出，我們可以透過以下方式儲存它：
class MyModule ( nn . Module ):
def __init__ ( self , sizes ):
super (). __init__ ()
self . layers = nn . ModuleList ([ nn . Linear ( in_f , out_f ) for in_f , out_f in zip ( sizes , sizes [ 1 :])])
self . trace = []
def forward ( self , x ):
for layer in self . layers :
x = layer ( x )
self . trace . append ( x )
return x model = MyModule ([ 1 , 16 , 32 ])
import torch
model ( torch . rand (( 4 , 1 )))
[ print ( trace . shape ) for trace in model . trace ] torch.Size([4, 16])
torch.Size([4, 32])
[None, None]
如果我們想在conv_block中切換到LearkyRelu該怎麼辦？我們可以使用ModuleDict建立Module的字典，並在需要時動態切換Module
def conv_block ( in_f , out_f , activation = 'relu' , * args , ** kwargs ):
activations = nn . ModuleDict ([
[ 'lrelu' , nn . LeakyReLU ()],
[ 'relu' , nn . ReLU ()]
])
return nn . Sequential (
nn . Conv2d ( in_f , out_f , * args , ** kwargs ),
nn . BatchNorm2d ( out_f ),
activations [ activation ]
) print ( conv_block ( 1 , 32 , 'lrelu' , kernel_size = 3 , padding = 1 ))
print ( conv_block ( 1 , 32 , 'relu' , kernel_size = 3 , padding = 1 )) Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): LeakyReLU(negative_slope=0.01)
)
Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
)
讓我們把一切都包起來吧！
def conv_block ( in_f , out_f , activation = 'relu' , * args , ** kwargs ):
activations = nn . ModuleDict ([
[ 'lrelu' , nn . LeakyReLU ()],
[ 'relu' , nn . ReLU ()]
])
return nn . Sequential (
nn . Conv2d ( in_f , out_f , * args , ** kwargs ),
nn . BatchNorm2d ( out_f ),
activations [ activation ]
)
def dec_block ( in_f , out_f ):
return nn . Sequential (
nn . Linear ( in_f , out_f ),
nn . Sigmoid ()
)
class MyEncoder ( nn . Module ):
def __init__ ( self , enc_sizes , * args , ** kwargs ):
super (). __init__ ()
self . conv_blocks = nn . Sequential ( * [ conv_block ( in_f , out_f , kernel_size = 3 , padding = 1 , * args , ** kwargs )
for in_f , out_f in zip ( enc_sizes , enc_sizes [ 1 :])])
def forward ( self , x ):
return self . conv_blocks ( x )
class MyDecoder ( nn . Module ):
def __init__ ( self , dec_sizes , n_classes ):
super (). __init__ ()
self . dec_blocks = nn . Sequential ( * [ dec_block ( in_f , out_f )
for in_f , out_f in zip ( dec_sizes , dec_sizes [ 1 :])])
self . last = nn . Linear ( dec_sizes [ - 1 ], n_classes )
def forward ( self , x ):
return self . dec_blocks ()
class MyCNNClassifier ( nn . Module ):
def __init__ ( self , in_c , enc_sizes , dec_sizes , n_classes , activation = 'relu' ):
super (). __init__ ()
self . enc_sizes = [ in_c , * enc_sizes ]
self . dec_sizes = [ 32 * 28 * 28 , * dec_sizes ]
self . encoder = MyEncoder ( self . enc_sizes , activation = activation )
self . decoder = MyDecoder ( dec_sizes , n_classes )
def forward ( self , x ):
x = self . encoder ( x )
x = x . flatten ( 1 ) # flat
x = self . decoder ( x )
return x model = MyCNNClassifier ( 1 , [ 32 , 64 ], [ 1024 , 512 ], 10 , activation = 'lrelu' )
print ( model ) MyCNNClassifier(
(encoder): MyEncoder(
(conv_blocks): Sequential(
(0): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): LeakyReLU(negative_slope=0.01)
)
(1): Sequential(
(0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): LeakyReLU(negative_slope=0.01)
)
)
)
(decoder): MyDecoder(
(dec_blocks): Sequential(
(0): Sequential(
(0): Linear(in_features=1024, out_features=512, bias=True)
(1): Sigmoid()
)
)
(last): Linear(in_features=512, out_features=10, bias=True)
)
)
所以，總而言之。
ModuleSequentialModuleListModuleDict這就是大家！
感謝您的閱讀
