class OurModule(nn.Module):
    def __init__(self, num_inputs, num_classes, dropout_prob=0.3):  
        super(OurModule, self).__init__()
        self.pipe = nn.Sequential(
            nn.Linear(num_inputs, 5),
            nn.ReLU(),
            nn.Linear(5, 20),
            nn.ReLU(),
            nn.Linear(20, num_classes),
            nn.Dropout(p=dropout_prob),
            nn.Softmax(dim=1)
        )

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

上面參數(shù)主要介紹下dropout_prob，模型訓(xùn)練時應(yīng)用Dropout的流程，概況一下描述就是：
1.隨機(jī)概率p隨機(jī)dropout部分神經(jīng)元，并前向傳播
2.計(jì)算前向傳播的損失，應(yīng)用反向傳播和梯度更新(對剩余的未被dropout的神經(jīng)元)
3.恢復(fù)所有神經(jīng)元的，并重復(fù)過程1

if __name__ == "__main__":
    net = OurModule(num_inputs=2, num_classes=3)
    print(net)
    v = torch.FloatTensor([[2, 3]])
    out = net(v)
    print(out)
    print("Cuda's availability is %s" % torch.cuda.is_available())
    if torch.cuda.is_available():
        print("Data from cuda: %s" % out.to('cuda'))

輸入?yún)?shù)為Tensor[2,3]，經(jīng)過三層NN，最后輸出softmax。

import torch as t
from torch import nn
from torch.nn import functional as F

# 假定輸入的圖像形狀為[3,64,64]
x = t.randn(10, 3, 64, 64)      # 10張 3個channel 大小為64x64的圖片

x = nn.Conv2d(3, 64, kernel_size=3, stride=3, padding=0)(x)
print(x.shape)


# 之前的特征圖尺寸為多少，只要設(shè)置為(1,1)，那么最終特征圖大小都為(1,1) 
# x = F.adaptive_avg_pool2d(x, [1,1])    # [b, 64, h, w] => [b, 64, 1, 1]
# print(x.shape)

# 將四維張量轉(zhuǎn)換為二維張量之后，才能作為全連接層的輸入
x = x.view(x.size(0), -1)  #view()的作用相當(dāng)于numpy中的reshape，重新定義矩陣的形狀。
print(x.shape)

# in_features由輸入張量的形狀決定，out_features則決定了輸出張量的形狀 
connected_layer = nn.Linear(in_features = 64*21*21, out_features = 10)

# 調(diào)用全連接層
output = connected_layer(x) 
print(output.shape)