File size: 1,021 Bytes
93db8d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | from torch import nn
class MessageFunction(nn.Module):
"""
Module which computes the message for a given interaction.
"""
def compute_message(self, raw_messages):
return None
class MLPMessageFunction(MessageFunction):
def __init__(self, raw_message_dimension, message_dimension):
super(MLPMessageFunction, self).__init__()
self.mlp = self.layers = nn.Sequential(
nn.Linear(raw_message_dimension, raw_message_dimension // 2),
nn.ReLU(),
nn.Linear(raw_message_dimension // 2, message_dimension),
)
def compute_message(self, raw_messages):
messages = self.mlp(raw_messages)
return messages
class IdentityMessageFunction(MessageFunction):
def compute_message(self, raw_messages):
return raw_messages
def get_message_function(module_type, raw_message_dimension, message_dimension):
if module_type == "mlp":
return MLPMessageFunction(raw_message_dimension, message_dimension)
elif module_type == "identity":
return IdentityMessageFunction()
|