| 43 | print("L_Decom ...") |
| 44 | |
| 45 | def forward(self, x, edge_indexs=None, edge_weights=None, y=None): |
| 46 | |
| 47 | x, x_m, x_s = self.decompose(x) |
| 48 | |
| 49 | |
| 50 | x1 = self.l_x1(x) |
| 51 | x1 = F.relu(x1) |
| 52 | x2 = self.l_x2(x1) |
| 53 | |
| 54 | m1 = self.l_m1(x_m) |
| 55 | m1 = F.relu(m1) |
| 56 | m2 = self.l_m2(m1) |
| 57 | |
| 58 | s1 = self.l_s1(x_s) |
| 59 | s1 = F.relu(s1) |
| 60 | s2 = self.l_s2(s1) |
| 61 | |
| 62 | out = x2 * s2 |
| 63 | |
| 64 | out = self.l_o1(out) |
| 65 | out = F.relu(out) |
| 66 | out = self.l_o2(out) + m2 |
| 67 | |
| 68 | if y is not None: |
| 69 | |
| 70 | _y, y_m, y_s = self.decompose(y) |
| 71 | |
| 72 | loss = self.L1(out, y) + 0.1*(self.L2(x2,_y) + self.L2(m2, y_m) + self.L2(s2, y_s)) |
| 73 | |
| 74 | return out, loss |
| 75 | |
| 76 | return out |
| 77 | |
| 78 | def decompose(self, y): |
| 79 | means = y.mean(-1, keepdim=True).detach() |