|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
|
|
|
|
class Bilinear(nn.Module): |
|
""" |
|
使用版本 |
|
A bilinear module that deals with broadcasting for efficient memory usage. |
|
Input: tensors of sizes (N x L1 x D1) and (N x L2 x D2) |
|
Output: tensor of size (N x L1 x L2 x O)""" |
|
|
|
def __init__(self, input1_size, input2_size, output_size, bias=True): |
|
super(Bilinear, self).__init__() |
|
|
|
self.input1_size = input1_size |
|
self.input2_size = input2_size |
|
self.output_size = output_size |
|
|
|
self.weight = nn.Parameter(torch.Tensor(input1_size, input2_size, output_size)) |
|
self.bias = nn.Parameter(torch.Tensor(output_size)) if bias else None |
|
|
|
self.reset_parameters() |
|
|
|
def reset_parameters(self): |
|
nn.init.zeros_(self.weight) |
|
|
|
def forward(self, input1, input2): |
|
input1_size = list(input1.size()) |
|
input2_size = list(input2.size()) |
|
|
|
intermediate = torch.mm(input1.view(-1, input1_size[-1]), self.weight.view(-1, self.input2_size * self.output_size),) |
|
|
|
input2 = input2.transpose(1, 2) |
|
output = intermediate.view(input1_size[0], input1_size[1] * self.output_size, input2_size[2]).bmm(input2) |
|
|
|
output = output.view(input1_size[0], input1_size[1], self.output_size, input2_size[1]).transpose(2, 3) |
|
|
|
if self.bias is not None: |
|
output = output + self.bias |
|
|
|
return output |
|
|
