| 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 | |