class CellTrack_Model(nn.Module): def __init__(self, hand_NodeEncoder_dic={}, learned_NodeEncoder_dic={}, intialize_EdgeEncoder_dic={}, message_passing={}, edge_classifier_dic={} ): super(CellTrack_Model, self).__init__() self.distance = CosineSimilarity() self.handcrafted_node_embedding = MLP(**hand_NodeEncoder_dic) self.learned_node_embedding = MLP(**learned_NodeEncoder_dic) self.learned_edge_embedding = MLP(**intialize_EdgeEncoder_dic) edge_mpnn_class = getattr(edge_mpnn, message_passing.target) self.message_passing = edge_mpnn_class(**message_passing.kwargs) self.edge_classifier = MLP(**edge_classifier_dic) def forward(self, x, edge_index, edge_feat): x1, x2 = x x_init = torch.cat((x1, x2), dim=-1) src, trg = edge_index similarity1 = self.distance(x_init[src], x_init[trg]) abs_init = torch.abs(x_init[src] - x_init[trg]) x1 = self.handcrafted_node_embedding(x1) x2 = self.learned_node_embedding(x2) x = torch.cat((x1, x2), dim=-1) src, trg = edge_index similarity2 = self.distance(x[src], x[trg]) edge_feat_in = torch.cat((abs_init, similarity1[:, None], x[src], x[trg], torch.abs(x[src] - x[trg]), similarity2[:, None]), dim=-1) edge_init_features = self.learned_edge_embedding(edge_feat_in) edge_feat_mp = self.message_passing(x, edge_index, edge_init_features) pred = self.edge_classifier(edge_feat_mp).squeeze() return pred

这段代码定义了一个名为 `CellTrack_Model` 的神经网络模型，该模型用于细胞轨迹跟踪任务。

在 `__init__` 方法中，模型的各个组件和参数被定义：

- `hand_NodeEncoder_dic`、`learned_NodeEncoder_dic`、`intialize_EdgeEncoder_dic`、`message_passing` 和 `edge_classifier_dic` 分别表示手工设计的节点编码器、学习得到的节点编码器、初始化的边编码器、消息传递参数和边分类器的参数字典。

在 `forward` 方法中，定义了模型的前向传播过程：

1. 首先对输入的节点特征 x 进行拆分，得到 x1 和 x2。然后将它们拼接成一个新的输入 x_init。

2. 计算 x_init 中源节点和目标节点之间的相似度 similarity1，以及它们的绝对差值 abs_init。

3. 分别通过手工设计的节点编码器和学习得到的节点编码器对 x1 和 x2 进行编码。

4. 将编码后的节点特征拼接成新的节点特征 x，并计算 x 中源节点和目标节点之间的相似度 similarity2。

5. 根据源节点和目标节点的特征，以及相似度和差值等特征，拼接成输入边特征 edge_feat_in。

6. 使用初始化的边编码器对输入边特征进行编码得到 edge_init_features。

7. 使用消息传递过程（message_passing）对节点特征和边特征进行传递和更新。

8. 将更新后的边特征通过边分类器进行预测，得到预测结果 pred。

9. 返回预测结果 pred。

这段代码实现了一个用于细胞轨迹跟踪的神经网络模型。通过对输入的节点特征进行编码和传递，最终预测边的标签。