從程式碼角度學習理解Pytorch學習框架03： 神經網路模組nn.Module的瞭解。

# coding=utf-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable


"""Pytorch中神經網路模組化介面nn的瞭解"""
"""
torch.nn是專門為神經網路設計的模組化介面。nn構建於autograd之上，可以用來定義和執行神經網路。
nn.Module是nn中十分重要的類,包含網路各層的定義及forward方法。


定義自已的網路：
    需要繼承nn.Module類，並實現forward方法。
    一般把網路中具有可學習引數的層放在建構函式__init__()中，
    不具有可學習引數的層(如ReLU)可放在建構函式中，也可不放在建構函式中(而在forward中使用nn.functional來代替)
    
    只要在nn.Module的子類中定義了forward函式，backward函式就會被自動實現(利用Autograd)。
    在forward函式中可以使用任何Variable支援的函式，畢竟在整個pytorch構建的圖中，是Variable在流動。還可以使用
    if,for,print,log等python語法.
    
    注：Pytorch基於nn.Module構建的模型中，只支援mini-batch的Variable輸入方式，
    比如，只有一張輸入圖片，也需要變成 N x C x H x W 的形式：
    
    input_image = torch.FloatTensor(1, 28, 28)
    input_image = Variable(input_image)
    input_image = input_image.unsqueeze(0)   # 1 x 1 x 28 x 28

    
"""


class LeNet(nn.Module):
    def __init__(self):
        # nn.Module的子類函式必須在建構函式中執行父類的建構函式
        super(LeNet, self).__init__()   # 等價與nn.Module.__init__()

        # nn.Conv2d返回的是一個Conv2d class的一個物件，該類中包含forward函式的實現
        # 當呼叫self.conv1(input)的時候，就會呼叫該類的forward函式
        self.conv1 = nn.Conv2d(1, 6, (5, 5))   # output (N, C_{out}, H_{out}, W_{out})`
        self.conv2 = nn.Conv2d(6, 16, (5, 5))
        self.fc1 = nn.Linear(256, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))  # F.max_pool2d的返回值是一個Variable
        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))

        # 返回值也是一個Variable物件
        return x


def output_name_and_params(net):
    for name, parameters in net.named_parameters():
        print('name: {}, param: {}'.format(name, parameters))


if __name__ == '__main__':
    net = LeNet()
    print('net: {}'.format(net))
    params = net.parameters()   # generator object
    print('params: {}'.format(params))
    output_name_and_params(net)

    input_image = torch.FloatTensor(10, 1, 28, 28)

    # 和tensorflow不一樣，pytorch中模型的輸入是一個Variable，而且是Variable在圖中流動，不是Tensor。
    # 這可以從forward中每一步的執行結果可以看出
    input_image = Variable(input_image)

    output = net(input_image)
    print('output: {}'.format(output))
    print('output.size: {}'.format(output.size()))