定义一个可学习的参数，并初始化

想念@思恋

于 2020-04-26 22:05:57 发布

阅读量880

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch

本文链接：https://blog.youkuaiyun.com/tailonh/article/details/105778418

pytorch 专栏收录该内容

47 篇文章

订阅专栏

本文介绍了一个自定义的PyTorch模块NewMatrices，该模块用于特定形状的矩阵运算，强调了初始化的重要性，并详细解释了其工作原理，包括权重和偏置的初始化策略。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

一定不要忘记初始化，否则可能会出错！！！

class NewMatrices(nn.Module):#[16, 1, 400]-->batch_size必须是16，若不满足，则丢弃
    def __init__(self, batch_size, input_dim, output_dim, bias_dim=128, dropout=0.1):
        super(NewMatrices, self).__init__()
        # self.dropout = nn.Dropout(dropout)
        self.weight = nn.Parameter(torch.Tensor(batch_size, input_dim, output_dim))
        self.bias = nn.Parameter(torch.Tensor(bias_dim))
        self.reset_parameters()
    def reset_parameters(self):
        torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in)
            torch.nn.init.uniform_(self.bias, -bound, bound)
    def forward(self, x):
        # x = self.dropout(x)
        return torch.bmm(self.weight, x)