pytorch1.0 实现多个层权重共享

原创

已于 2022-04-25 19:03:33 修改 · 7k 阅读

5 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #权重共享 #动态网络 #python

于 2019-08-08 11:35:26 首次发布

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