| 74 | |
| 75 | @R.register("models.FusionNetwork") |
| 76 | class FusionNetwork(nn.Module, core.Configurable): |
| 77 | |
| 78 | def __init__(self, sequence_model, structure_model): |
| 79 | super(FusionNetwork, self).__init__() |
| 80 | self.sequence_model = sequence_model |
| 81 | self.structure_model = structure_model |
| 82 | self.output_dim = structure_model.output_dim |
| 83 | |
| 84 | def forward(self, graph, input, all_loss=None, metric=None, batch=None): |
| 85 | # Sequence model |
| 86 | output1 = self.sequence_model(graph, input, all_loss, metric) |
| 87 | if "logits" in output1: |
| 88 | logits = output1["logits"] |
| 89 | else: |
| 90 | logits = None |
| 91 | node_output1 = output1["residue_feature"] |
| 92 | |
| 93 | # Structure model |
| 94 | output2 = self.structure_model(graph, node_output1, all_loss, metric) |
| 95 | node_feature = output2["node_feature"] |
| 96 | graph_feature = output2["graph_feature"] |
| 97 | |
| 98 | return { |
| 99 | "graph_feature": graph_feature, |
| 100 | "node_feature": node_feature, |
| 101 | "sequence_logits": logits, |
| 102 | } |
nothing calls this directly
no outgoing calls
no test coverage detected