Spaces:

ashuonnet
/

tgn

Runtime error

Upload tgn.py

eb859de verified over 1 year ago

14.6 kB

import logging
import numpy as np
import torch
from collections import defaultdict

from utils.utils import MergeLayer
from modules.memory import Memory
from modules.message_aggregator import get_message_aggregator
from modules.message_function import get_message_function
from modules.memory_updater import get_memory_updater
from modules.embedding_module import get_embedding_module
from model.time_encoding import TimeEncode


class TGN(torch.nn.Module):
  def __init__(self, neighbor_finder, node_features, edge_features, device, n_layers=2,
               n_heads=2, dropout=0.1, use_memory=False,
               memory_update_at_start=True, message_dimension=100,
               memory_dimension=500, embedding_module_type="graph_attention",
               message_function="mlp",
               mean_time_shift_src=0, std_time_shift_src=1, mean_time_shift_dst=0,
               std_time_shift_dst=1, n_neighbors=None, aggregator_type="last",
               memory_updater_type="gru",
               use_destination_embedding_in_message=False,
               use_source_embedding_in_message=False,
               dyrep=False):
    super(TGN, self).__init__()

    self.n_layers = n_layers
    self.neighbor_finder = neighbor_finder
    self.device = device
    self.logger = logging.getLogger(__name__)

    self.node_raw_features = torch.from_numpy(node_features.astype(np.float32)).to(device)
    self.edge_raw_features = torch.from_numpy(edge_features.astype(np.float32)).to(device)

    self.n_node_features = self.node_raw_features.shape[1]
    self.n_nodes = self.node_raw_features.shape[0]
    self.n_edge_features = self.edge_raw_features.shape[1]
    self.embedding_dimension = self.n_node_features
    self.n_neighbors = n_neighbors
    self.embedding_module_type = embedding_module_type
    self.use_destination_embedding_in_message = use_destination_embedding_in_message
    self.use_source_embedding_in_message = use_source_embedding_in_message
    self.dyrep = dyrep

    self.use_memory = use_memory
    self.time_encoder = TimeEncode(dimension=self.n_node_features)
    self.memory = None

    self.mean_time_shift_src = mean_time_shift_src
    self.std_time_shift_src = std_time_shift_src
    self.mean_time_shift_dst = mean_time_shift_dst
    self.std_time_shift_dst = std_time_shift_dst

    if self.use_memory:
      self.memory_dimension = memory_dimension
      self.memory_update_at_start = memory_update_at_start
      raw_message_dimension = 2 * self.memory_dimension + self.n_edge_features + \
                              self.time_encoder.dimension
      message_dimension = message_dimension if message_function != "identity" else raw_message_dimension
      self.memory = Memory(n_nodes=self.n_nodes,
                           memory_dimension=self.memory_dimension,
                           input_dimension=message_dimension,
                           message_dimension=message_dimension,
                           device=device)
      self.message_aggregator = get_message_aggregator(aggregator_type=aggregator_type,
                                                       device=device)
      self.message_function = get_message_function(module_type=message_function,
                                                   raw_message_dimension=raw_message_dimension,
                                                   message_dimension=message_dimension)
      self.memory_updater = get_memory_updater(module_type=memory_updater_type,
                                               memory=self.memory,
                                               message_dimension=message_dimension,
                                               memory_dimension=self.memory_dimension,
                                               device=device)

    self.embedding_module_type = embedding_module_type

    self.embedding_module = get_embedding_module(module_type=embedding_module_type,
                                                 node_features=self.node_raw_features,
                                                 edge_features=self.edge_raw_features,
                                                 memory=self.memory,
                                                 neighbor_finder=self.neighbor_finder,
                                                 time_encoder=self.time_encoder,
                                                 n_layers=self.n_layers,
                                                 n_node_features=self.n_node_features,
                                                 n_edge_features=self.n_edge_features,
                                                 n_time_features=self.n_node_features,
                                                 embedding_dimension=self.embedding_dimension,
                                                 device=self.device,
                                                 n_heads=n_heads, dropout=dropout,
                                                 use_memory=use_memory,
                                                 n_neighbors=self.n_neighbors)

    # MLP to compute probability on an edge given two node embeddings
    self.affinity_score = MergeLayer(self.n_node_features, self.n_node_features,
                                     self.n_node_features,
                                     1)

  def compute_temporal_embeddings(self, source_nodes, destination_nodes, negative_nodes, edge_times,
                                  edge_idxs, n_neighbors=20):
    """
    Compute temporal embeddings for sources, destinations, and negatively sampled destinations.
    source_nodes [batch_size]: source ids.
    :param destination_nodes [batch_size]: destination ids
    :param negative_nodes [batch_size]: ids of negative sampled destination
    :param edge_times [batch_size]: timestamp of interaction
    :param edge_idxs [batch_size]: index of interaction
    :param n_neighbors [scalar]: number of temporal neighbor to consider in each convolutional
    layer
    :return: Temporal embeddings for sources, destinations and negatives
    """

    n_samples = len(source_nodes)
    nodes = np.concatenate([source_nodes, destination_nodes, negative_nodes])
    positives = np.concatenate([source_nodes, destination_nodes])
    timestamps = np.concatenate([edge_times, edge_times, edge_times])

    memory = None
    time_diffs = None
    if self.use_memory:
      if self.memory_update_at_start:
        # Update memory for all nodes with messages stored in previous batches
        memory, last_update = self.get_updated_memory(list(range(self.n_nodes)),
                                                      self.memory.messages)
      else:
        memory = self.memory.get_memory(list(range(self.n_nodes)))
        last_update = self.memory.last_update

      ### Compute differences between the time the memory of a node was last updated,
      ### and the time for which we want to compute the embedding of a node
      source_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[
        source_nodes].long()
      source_time_diffs = (source_time_diffs - self.mean_time_shift_src) / self.std_time_shift_src
      destination_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[
        destination_nodes].long()
      destination_time_diffs = (destination_time_diffs - self.mean_time_shift_dst) / self.std_time_shift_dst
      negative_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[
        negative_nodes].long()
      negative_time_diffs = (negative_time_diffs - self.mean_time_shift_dst) / self.std_time_shift_dst

      time_diffs = torch.cat([source_time_diffs, destination_time_diffs, negative_time_diffs],
                             dim=0)

    # Compute the embeddings using the embedding module
    node_embedding = self.embedding_module.compute_embedding(memory=memory,
                                                             source_nodes=nodes,
                                                             timestamps=timestamps,
                                                             n_layers=self.n_layers,
                                                             n_neighbors=n_neighbors,
                                                             time_diffs=time_diffs)

    source_node_embedding = node_embedding[:n_samples]
    destination_node_embedding = node_embedding[n_samples: 2 * n_samples]
    negative_node_embedding = node_embedding[2 * n_samples:]

    if self.use_memory:
      if self.memory_update_at_start:
        # Persist the updates to the memory only for sources and destinations (since now we have
        # new messages for them)
        self.update_memory(positives, self.memory.messages)

        assert torch.allclose(memory[positives], self.memory.get_memory(positives), atol=1e-5), \
          "Something wrong in how the memory was updated"

        # Remove messages for the positives since we have already updated the memory using them
        self.memory.clear_messages(positives)

      unique_sources, source_id_to_messages = self.get_raw_messages(source_nodes,
                                                                    source_node_embedding,
                                                                    destination_nodes,
                                                                    destination_node_embedding,
                                                                    edge_times, edge_idxs)
      unique_destinations, destination_id_to_messages = self.get_raw_messages(destination_nodes,
                                                                              destination_node_embedding,
                                                                              source_nodes,
                                                                              source_node_embedding,
                                                                              edge_times, edge_idxs)
      if self.memory_update_at_start:
        self.memory.store_raw_messages(unique_sources, source_id_to_messages)
        self.memory.store_raw_messages(unique_destinations, destination_id_to_messages)
      else:
        self.update_memory(unique_sources, source_id_to_messages)
        self.update_memory(unique_destinations, destination_id_to_messages)

      if self.dyrep:
        source_node_embedding = memory[source_nodes]
        destination_node_embedding = memory[destination_nodes]
        negative_node_embedding = memory[negative_nodes]

    return source_node_embedding, destination_node_embedding, negative_node_embedding

  def compute_edge_probabilities(self, source_nodes, destination_nodes, negative_nodes, edge_times,
                                 edge_idxs, n_neighbors=20):
    """
    Compute probabilities for edges between sources and destination and between sources and
    negatives by first computing temporal embeddings using the TGN encoder and then feeding them
    into the MLP decoder.
    :param destination_nodes [batch_size]: destination ids
    :param negative_nodes [batch_size]: ids of negative sampled destination
    :param edge_times [batch_size]: timestamp of interaction
    :param edge_idxs [batch_size]: index of interaction
    :param n_neighbors [scalar]: number of temporal neighbor to consider in each convolutional
    layer
    :return: Probabilities for both the positive and negative edges
    """
    n_samples = len(source_nodes)
    source_node_embedding, destination_node_embedding, negative_node_embedding = self.compute_temporal_embeddings(
      source_nodes, destination_nodes, negative_nodes, edge_times, edge_idxs, n_neighbors)

    score = self.affinity_score(torch.cat([source_node_embedding, source_node_embedding], dim=0),
                                torch.cat([destination_node_embedding,
                                           negative_node_embedding])).squeeze(dim=0)
    pos_score = score[:n_samples]
    neg_score = score[n_samples:]

    return pos_score.sigmoid(), neg_score.sigmoid()

  def update_memory(self, nodes, messages):
    # Aggregate messages for the same nodes
    unique_nodes, unique_messages, unique_timestamps = \
      self.message_aggregator.aggregate(
        nodes,
        messages)

    if len(unique_nodes) > 0:
      unique_messages = self.message_function.compute_message(unique_messages)

    # Update the memory with the aggregated messages
    self.memory_updater.update_memory(unique_nodes, unique_messages,
                                      timestamps=unique_timestamps)

  def get_updated_memory(self, nodes, messages):
    # Aggregate messages for the same nodes
    unique_nodes, unique_messages, unique_timestamps = \
      self.message_aggregator.aggregate(
        nodes,
        messages)

    if len(unique_nodes) > 0:
      unique_messages = self.message_function.compute_message(unique_messages)

    updated_memory, updated_last_update = self.memory_updater.get_updated_memory(unique_nodes,
                                                                                 unique_messages,
                                                                                 timestamps=unique_timestamps)

    return updated_memory, updated_last_update

  def get_raw_messages(self, source_nodes, source_node_embedding, destination_nodes,
                       destination_node_embedding, edge_times, edge_idxs):
    edge_times = torch.from_numpy(edge_times).float().to(self.device)
    edge_features = self.edge_raw_features[edge_idxs]

    source_memory = self.memory.get_memory(source_nodes) if not \
      self.use_source_embedding_in_message else source_node_embedding
    destination_memory = self.memory.get_memory(destination_nodes) if \
      not self.use_destination_embedding_in_message else destination_node_embedding

    source_time_delta = edge_times - self.memory.last_update[source_nodes]
    source_time_delta_encoding = self.time_encoder(source_time_delta.unsqueeze(dim=1)).view(len(
      source_nodes), -1)

    source_message = torch.cat([source_memory, destination_memory, edge_features,
                                source_time_delta_encoding],
                               dim=1)
    messages = defaultdict(list)
    unique_sources = np.unique(source_nodes)

    for i in range(len(source_nodes)):
      messages[source_nodes[i]].append((source_message[i], edge_times[i]))

    return unique_sources, messages

  def set_neighbor_finder(self, neighbor_finder):
    self.neighbor_finder = neighbor_finder
    self.embedding_module.neighbor_finder = neighbor_finder

	import logging
	import numpy as np
	import torch
	from collections import defaultdict

	from utils.utils import MergeLayer
	from modules.memory import Memory
	from modules.message_aggregator import get_message_aggregator
	from modules.message_function import get_message_function
	from modules.memory_updater import get_memory_updater
	from modules.embedding_module import get_embedding_module
	from model.time_encoding import TimeEncode


	class TGN(torch.nn.Module):
	def __init__(self, neighbor_finder, node_features, edge_features, device, n_layers=2,
	n_heads=2, dropout=0.1, use_memory=False,
	memory_update_at_start=True, message_dimension=100,
	memory_dimension=500, embedding_module_type="graph_attention",
	message_function="mlp",
	mean_time_shift_src=0, std_time_shift_src=1, mean_time_shift_dst=0,
	std_time_shift_dst=1, n_neighbors=None, aggregator_type="last",
	memory_updater_type="gru",
	use_destination_embedding_in_message=False,
	use_source_embedding_in_message=False,
	dyrep=False):
	super(TGN, self).__init__()

	self.n_layers = n_layers
	self.neighbor_finder = neighbor_finder
	self.device = device
	self.logger = logging.getLogger(__name__)

	self.node_raw_features = torch.from_numpy(node_features.astype(np.float32)).to(device)
	self.edge_raw_features = torch.from_numpy(edge_features.astype(np.float32)).to(device)

	self.n_node_features = self.node_raw_features.shape[1]
	self.n_nodes = self.node_raw_features.shape[0]
	self.n_edge_features = self.edge_raw_features.shape[1]
	self.embedding_dimension = self.n_node_features
	self.n_neighbors = n_neighbors
	self.embedding_module_type = embedding_module_type
	self.use_destination_embedding_in_message = use_destination_embedding_in_message
	self.use_source_embedding_in_message = use_source_embedding_in_message
	self.dyrep = dyrep

	self.use_memory = use_memory
	self.time_encoder = TimeEncode(dimension=self.n_node_features)
	self.memory = None

	self.mean_time_shift_src = mean_time_shift_src
	self.std_time_shift_src = std_time_shift_src
	self.mean_time_shift_dst = mean_time_shift_dst
	self.std_time_shift_dst = std_time_shift_dst

	if self.use_memory:
	self.memory_dimension = memory_dimension
	self.memory_update_at_start = memory_update_at_start
	raw_message_dimension = 2 * self.memory_dimension + self.n_edge_features + \
	self.time_encoder.dimension
	message_dimension = message_dimension if message_function != "identity" else raw_message_dimension
	self.memory = Memory(n_nodes=self.n_nodes,
	memory_dimension=self.memory_dimension,
	input_dimension=message_dimension,
	message_dimension=message_dimension,
	device=device)
	self.message_aggregator = get_message_aggregator(aggregator_type=aggregator_type,
	device=device)
	self.message_function = get_message_function(module_type=message_function,
	raw_message_dimension=raw_message_dimension,
	message_dimension=message_dimension)
	self.memory_updater = get_memory_updater(module_type=memory_updater_type,
	memory=self.memory,
	message_dimension=message_dimension,
	memory_dimension=self.memory_dimension,
	device=device)

	self.embedding_module_type = embedding_module_type

	self.embedding_module = get_embedding_module(module_type=embedding_module_type,
	node_features=self.node_raw_features,
	edge_features=self.edge_raw_features,
	memory=self.memory,
	neighbor_finder=self.neighbor_finder,
	time_encoder=self.time_encoder,
	n_layers=self.n_layers,
	n_node_features=self.n_node_features,
	n_edge_features=self.n_edge_features,
	n_time_features=self.n_node_features,
	embedding_dimension=self.embedding_dimension,
	device=self.device,
	n_heads=n_heads, dropout=dropout,
	use_memory=use_memory,
	n_neighbors=self.n_neighbors)

	# MLP to compute probability on an edge given two node embeddings
	self.affinity_score = MergeLayer(self.n_node_features, self.n_node_features,
	self.n_node_features,
	1)

	def compute_temporal_embeddings(self, source_nodes, destination_nodes, negative_nodes, edge_times,
	edge_idxs, n_neighbors=20):
	"""
	Compute temporal embeddings for sources, destinations, and negatively sampled destinations.

	source_nodes [batch_size]: source ids.
	:param destination_nodes [batch_size]: destination ids
	:param negative_nodes [batch_size]: ids of negative sampled destination
	:param edge_times [batch_size]: timestamp of interaction
	:param edge_idxs [batch_size]: index of interaction
	:param n_neighbors [scalar]: number of temporal neighbor to consider in each convolutional
	layer
	:return: Temporal embeddings for sources, destinations and negatives
	"""

	n_samples = len(source_nodes)
	nodes = np.concatenate([source_nodes, destination_nodes, negative_nodes])
	positives = np.concatenate([source_nodes, destination_nodes])
	timestamps = np.concatenate([edge_times, edge_times, edge_times])

	memory = None
	time_diffs = None
	if self.use_memory:
	if self.memory_update_at_start:
	# Update memory for all nodes with messages stored in previous batches
	memory, last_update = self.get_updated_memory(list(range(self.n_nodes)),
	self.memory.messages)
	else:
	memory = self.memory.get_memory(list(range(self.n_nodes)))
	last_update = self.memory.last_update

	### Compute differences between the time the memory of a node was last updated,
	### and the time for which we want to compute the embedding of a node
	source_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[
	source_nodes].long()
	source_time_diffs = (source_time_diffs - self.mean_time_shift_src) / self.std_time_shift_src
	destination_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[
	destination_nodes].long()
	destination_time_diffs = (destination_time_diffs - self.mean_time_shift_dst) / self.std_time_shift_dst
	negative_time_diffs = torch.LongTensor(edge_times).to(self.device) - last_update[
	negative_nodes].long()
	negative_time_diffs = (negative_time_diffs - self.mean_time_shift_dst) / self.std_time_shift_dst

	time_diffs = torch.cat([source_time_diffs, destination_time_diffs, negative_time_diffs],
	dim=0)

	# Compute the embeddings using the embedding module
	node_embedding = self.embedding_module.compute_embedding(memory=memory,
	source_nodes=nodes,
	timestamps=timestamps,
	n_layers=self.n_layers,
	n_neighbors=n_neighbors,
	time_diffs=time_diffs)

	source_node_embedding = node_embedding[:n_samples]
	destination_node_embedding = node_embedding[n_samples: 2 * n_samples]
	negative_node_embedding = node_embedding[2 * n_samples:]

	if self.use_memory:
	if self.memory_update_at_start:
	# Persist the updates to the memory only for sources and destinations (since now we have
	# new messages for them)
	self.update_memory(positives, self.memory.messages)

	assert torch.allclose(memory[positives], self.memory.get_memory(positives), atol=1e-5), \
	"Something wrong in how the memory was updated"

	# Remove messages for the positives since we have already updated the memory using them
	self.memory.clear_messages(positives)

	unique_sources, source_id_to_messages = self.get_raw_messages(source_nodes,
	source_node_embedding,
	destination_nodes,
	destination_node_embedding,
	edge_times, edge_idxs)
	unique_destinations, destination_id_to_messages = self.get_raw_messages(destination_nodes,
	destination_node_embedding,
	source_nodes,
	source_node_embedding,
	edge_times, edge_idxs)
	if self.memory_update_at_start:
	self.memory.store_raw_messages(unique_sources, source_id_to_messages)
	self.memory.store_raw_messages(unique_destinations, destination_id_to_messages)
	else:
	self.update_memory(unique_sources, source_id_to_messages)
	self.update_memory(unique_destinations, destination_id_to_messages)

	if self.dyrep:
	source_node_embedding = memory[source_nodes]
	destination_node_embedding = memory[destination_nodes]
	negative_node_embedding = memory[negative_nodes]

	return source_node_embedding, destination_node_embedding, negative_node_embedding

	def compute_edge_probabilities(self, source_nodes, destination_nodes, negative_nodes, edge_times,
	edge_idxs, n_neighbors=20):
	"""
	Compute probabilities for edges between sources and destination and between sources and
	negatives by first computing temporal embeddings using the TGN encoder and then feeding them
	into the MLP decoder.
	:param destination_nodes [batch_size]: destination ids
	:param negative_nodes [batch_size]: ids of negative sampled destination
	:param edge_times [batch_size]: timestamp of interaction
	:param edge_idxs [batch_size]: index of interaction
	:param n_neighbors [scalar]: number of temporal neighbor to consider in each convolutional
	layer
	:return: Probabilities for both the positive and negative edges
	"""
	n_samples = len(source_nodes)
	source_node_embedding, destination_node_embedding, negative_node_embedding = self.compute_temporal_embeddings(
	source_nodes, destination_nodes, negative_nodes, edge_times, edge_idxs, n_neighbors)

	score = self.affinity_score(torch.cat([source_node_embedding, source_node_embedding], dim=0),
	torch.cat([destination_node_embedding,
	negative_node_embedding])).squeeze(dim=0)
	pos_score = score[:n_samples]
	neg_score = score[n_samples:]

	return pos_score.sigmoid(), neg_score.sigmoid()

	def update_memory(self, nodes, messages):
	# Aggregate messages for the same nodes
	unique_nodes, unique_messages, unique_timestamps = \
	self.message_aggregator.aggregate(
	nodes,
	messages)

	if len(unique_nodes) > 0:
	unique_messages = self.message_function.compute_message(unique_messages)

	# Update the memory with the aggregated messages
	self.memory_updater.update_memory(unique_nodes, unique_messages,
	timestamps=unique_timestamps)

	def get_updated_memory(self, nodes, messages):
	# Aggregate messages for the same nodes
	unique_nodes, unique_messages, unique_timestamps = \
	self.message_aggregator.aggregate(
	nodes,
	messages)

	if len(unique_nodes) > 0:
	unique_messages = self.message_function.compute_message(unique_messages)

	updated_memory, updated_last_update = self.memory_updater.get_updated_memory(unique_nodes,
	unique_messages,
	timestamps=unique_timestamps)

	return updated_memory, updated_last_update

	def get_raw_messages(self, source_nodes, source_node_embedding, destination_nodes,
	destination_node_embedding, edge_times, edge_idxs):
	edge_times = torch.from_numpy(edge_times).float().to(self.device)
	edge_features = self.edge_raw_features[edge_idxs]

	source_memory = self.memory.get_memory(source_nodes) if not \
	self.use_source_embedding_in_message else source_node_embedding
	destination_memory = self.memory.get_memory(destination_nodes) if \
	not self.use_destination_embedding_in_message else destination_node_embedding

	source_time_delta = edge_times - self.memory.last_update[source_nodes]
	source_time_delta_encoding = self.time_encoder(source_time_delta.unsqueeze(dim=1)).view(len(
	source_nodes), -1)

	source_message = torch.cat([source_memory, destination_memory, edge_features,
	source_time_delta_encoding],
	dim=1)
	messages = defaultdict(list)
	unique_sources = np.unique(source_nodes)

	for i in range(len(source_nodes)):
	messages[source_nodes[i]].append((source_message[i], edge_times[i]))

	return unique_sources, messages

	def set_neighbor_finder(self, neighbor_finder):
	self.neighbor_finder = neighbor_finder
	self.embedding_module.neighbor_finder = neighbor_finder