Spaces:
Sleeping
Sleeping
| from collections import defaultdict | |
| import torch | |
| import numpy as np | |
| class MessageAggregator(torch.nn.Module): | |
| """ | |
| Abstract class for the message aggregator module, which given a batch of node ids and | |
| corresponding messages, aggregates messages with the same node id. | |
| """ | |
| def __init__(self, device): | |
| super(MessageAggregator, self).__init__() | |
| self.device = device | |
| def aggregate(self, node_ids, messages): | |
| """ | |
| Given a list of node ids, and a list of messages of the same length, aggregate different | |
| messages for the same id using one of the possible strategies. | |
| :param node_ids: A list of node ids of length batch_size | |
| :param messages: A tensor of shape [batch_size, message_length] | |
| :param timestamps A tensor of shape [batch_size] | |
| :return: A tensor of shape [n_unique_node_ids, message_length] with the aggregated messages | |
| """ | |
| def group_by_id(self, node_ids, messages, timestamps): | |
| node_id_to_messages = defaultdict(list) | |
| for i, node_id in enumerate(node_ids): | |
| node_id_to_messages[node_id].append((messages[i], timestamps[i])) | |
| return node_id_to_messages | |
| class LastMessageAggregator(MessageAggregator): | |
| def __init__(self, device): | |
| super(LastMessageAggregator, self).__init__(device) | |
| def aggregate(self, node_ids, messages): | |
| """Only keep the last message for each node""" | |
| unique_node_ids = np.unique(node_ids) | |
| unique_messages = [] | |
| unique_timestamps = [] | |
| to_update_node_ids = [] | |
| for node_id in unique_node_ids: | |
| if len(messages[node_id]) > 0: | |
| to_update_node_ids.append(node_id) | |
| unique_messages.append(messages[node_id][-1][0]) | |
| unique_timestamps.append(messages[node_id][-1][1]) | |
| unique_messages = torch.stack(unique_messages) if len(to_update_node_ids) > 0 else [] | |
| unique_timestamps = torch.stack(unique_timestamps) if len(to_update_node_ids) > 0 else [] | |
| return to_update_node_ids, unique_messages, unique_timestamps | |
| class MeanMessageAggregator(MessageAggregator): | |
| def __init__(self, device): | |
| super(MeanMessageAggregator, self).__init__(device) | |
| def aggregate(self, node_ids, messages): | |
| """Only keep the last message for each node""" | |
| unique_node_ids = np.unique(node_ids) | |
| unique_messages = [] | |
| unique_timestamps = [] | |
| to_update_node_ids = [] | |
| n_messages = 0 | |
| for node_id in unique_node_ids: | |
| if len(messages[node_id]) > 0: | |
| n_messages += len(messages[node_id]) | |
| to_update_node_ids.append(node_id) | |
| unique_messages.append(torch.mean(torch.stack([m[0] for m in messages[node_id]]), dim=0)) | |
| unique_timestamps.append(messages[node_id][-1][1]) | |
| unique_messages = torch.stack(unique_messages) if len(to_update_node_ids) > 0 else [] | |
| unique_timestamps = torch.stack(unique_timestamps) if len(to_update_node_ids) > 0 else [] | |
| return to_update_node_ids, unique_messages, unique_timestamps | |
| def get_message_aggregator(aggregator_type, device): | |
| if aggregator_type == "last": | |
| return LastMessageAggregator(device=device) | |
| elif aggregator_type == "mean": | |
| return MeanMessageAggregator(device=device) | |
| else: | |
| raise ValueError("Message aggregator {} not implemented".format(aggregator_type)) | |