pytorch1.0 實現多個層權重共享

1、在模型前向傳播時，可以多次重用同一個模塊實現權重共享。

2、用python循環語句或條件語句在每個前向傳播時構建一個動態計算圖，所以下面這個模型是一個動態網絡（動態控制流程）

import torch
import torch.nn as nn
import random
import matplotlib.pyplot as plt

# 繪製loss曲線
def plot_curve(data):
    fig = plt.figure()
    plt.plot(range(len(data)), data, color='blue')
    plt.legend(['value'], loc='upper right')
    plt.xlabel('step')
    plt.ylabel('value')
    plt.show()


class DynamicNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super(DynamicNet, self).__init__()
        self.input_linear = nn.Linear(D_in, H)
        self.middle_linear = nn.Linear(H, H)
        self.output_linear = nn.Linear(H, D_out)

    def forward(self, x):
        h_relu = self.input_linear(x).clamp(min=0)
        # 重複利用Middle linear模塊
        for _ in range(random.randint(0, 3)):
            h_relu = self.middle_linear(h_relu).clamp(min=0)
        y_pred = self.output_linear(h_relu)
        return y_pred


# N是批大小；D是輸入維度
# H是隱藏層維度；D_out是輸出維度
N, D_in, H, D_out = 64, 1000, 100, 10

# 模擬訓練數據
x = torch.randn(N, D_in)
y = torch.randn(N, D_out)

model = DynamicNet(D_in, H, D_out)
criterion = nn.MSELoss(reduction='sum')
# 用平凡的隨機梯度下降訓練這個奇怪的模型是困難的，所以我們使用了momentum方法。
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9)

loss_list = []
for t in range(500):
    # 前向傳播
    y_pred = model(x)
    # 計算損失
    loss = criterion(y_pred, y)
    loss_list.append(loss.item())
    # 清零梯度，反向傳播，更新權重
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

plot_curve(loss_list)