| 31 | return A |
| 32 | |
| 33 | class Attn_Modality_Gated(nn.Module): |
| 34 | # Adapted from https://github.com/mahmoodlab/PORPOISE |
| 35 | def __init__(self, gate_h1, gate_h2, gate_h3, dim1_og, dim2_og, dim3_og, use_bilinear=[True,True,True], scale=[1,1,1], p_dropout_fc=0.25): |
| 36 | super(Attn_Modality_Gated, self).__init__() |
| 37 | |
| 38 | self.gate_h1 = gate_h1 #[boolean] |
| 39 | self.gate_h2 = gate_h2 #[boolean] |
| 40 | self.gate_h3 = gate_h3 #[boolean] |
| 41 | self.use_bilinear = use_bilinear #[boolean] |
| 42 | |
| 43 | # can perform attention on latent vectors of lower dimension |
| 44 | dim1, dim2, dim3 = dim1_og//scale[0], dim2_og//scale[1], dim3_og//scale[2] |
| 45 | |
| 46 | # attention gate of each modality |
| 47 | if self.gate_h1: |
| 48 | self.linear_h1 = nn.Sequential(nn.Linear(dim1_og, dim1), nn.ReLU()) |
| 49 | self.linear_z1 = nn.Bilinear(dim1_og, dim2_og+dim3_og, dim1) if self.use_bilinear[0] else nn.Sequential(nn.Linear(dim1_og+dim2_og+dim3_og, dim1)) |
| 50 | self.linear_o1 = nn.Sequential(nn.Linear(dim1, dim1), nn.ReLU(), nn.Dropout(p=p_dropout_fc)) |
| 51 | else: |
| 52 | self.linear_h1, self.linear_o1 = nn.Identity(), nn.Identity() |
| 53 | |
| 54 | if self.gate_h2: |
| 55 | self.linear_h2 = nn.Sequential(nn.Linear(dim2_og, dim2), nn.ReLU()) |
| 56 | self.linear_z2 = nn.Bilinear(dim2_og, dim1_og+dim3_og, dim2) if self.use_bilinear[1] else nn.Sequential(nn.Linear(dim1_og+dim2_og+dim3_og, dim2)) |
| 57 | self.linear_o2 = nn.Sequential(nn.Linear(dim2, dim2), nn.ReLU(), nn.Dropout(p=p_dropout_fc)) |
| 58 | else: |
| 59 | self.linear_h2, self.linear_o2 = nn.Identity(), nn.Identity() |
| 60 | |
| 61 | if self.gate_h3: |
| 62 | self.linear_h3 = nn.Sequential(nn.Linear(dim3_og, dim3), nn.ReLU()) |
| 63 | self.linear_z3 = nn.Bilinear(dim3_og, dim1_og+dim2_og, dim3) if self.use_bilinear[2] else nn.Sequential(nn.Linear(dim1_og+dim2_og+dim3_og, dim3)) |
| 64 | self.linear_o3 = nn.Sequential(nn.Linear(dim3, dim3), nn.ReLU(), nn.Dropout(p=p_dropout_fc)) |
| 65 | else: |
| 66 | self.linear_h3, self.linear_o3 = nn.Identity(), nn.Identity() |
| 67 | |
| 68 | def forward(self, x1, x2, x3): |
| 69 | |
| 70 | if self.gate_h1: |
| 71 | h1 = self.linear_h1(x1) #breaks colli of h1 |
| 72 | z1 = self.linear_z1(x1, torch.cat([x2,x3], dim=-1)) if self.use_bilinear[0] else self.linear_z1(torch.cat((x1, x2, x3), dim=-1)) #creates a vector combining both modalities |
| 73 | o1 = self.linear_o1(nn.Sigmoid()(z1)*h1) #update modality input |
| 74 | else: |
| 75 | h1 = self.linear_h1(x1) |
| 76 | o1 = self.linear_o1(h1) |
| 77 | |
| 78 | if self.gate_h2: |
| 79 | h2 = self.linear_h2(x2) |
| 80 | z2 = self.linear_z2(x2, torch.cat([x1,x3], dim=-1)) if self.use_bilinear[1] else self.linear_z2(torch.cat((x1, x2, x3), dim=-1)) |
| 81 | o2 = self.linear_o2(nn.Sigmoid()(z2)*h2) |
| 82 | else: |
| 83 | h2 = self.linear_h2(x2) |
| 84 | o2 = self.linear_o2(h2) |
| 85 | |
| 86 | if self.gate_h3: |
| 87 | h3 = self.linear_h3(x3) |
| 88 | z3 = self.linear_z3(x3, torch.cat([x1,x2], dim=-1)) if self.use_bilinear[2] else self.linear_z3(torch.cat((x1, x2, x3), dim=-1)) |
| 89 | o3 = self.linear_o3(nn.Sigmoid()(z3)*h3) |
| 90 | else: |