MVGNN-PPIS / model /self_attention.py

Upload folder using huggingface_hub

ad9fbbf verified about 2 months ago

8.08 kB

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
import math
# The following gather functions
def gather_edges(edges, neighbor_idx):
    # Features [B,N,N,C] at Neighbor indices [B,N,K] => Neighbor features [B,N,K,C]
    neighbors = neighbor_idx.unsqueeze(-1).expand(-1, -1, -1, edges.size(-1))
    edge_features = torch.gather(edges, 2, neighbors)
    return edge_features

def gather_nodes(nodes, neighbor_idx):
    # Features [B,N,C] at Neighbor indices [B,N,K] => [B,N,K,C]
    # Flatten and expand indices per batch [B,N,K] => [B,NK] => [B,NK,C]
    neighbors_flat = neighbor_idx.view((neighbor_idx.shape[0], -1))
    neighbors_flat = neighbors_flat.unsqueeze(-1).expand(-1, -1, nodes.size(2))
    # Gather and re-pack
    neighbor_features = torch.gather(nodes, 1, neighbors_flat)
    neighbor_features = neighbor_features.view(list(neighbor_idx.shape)[:3] + [-1])
    return neighbor_features

def cat_neighbors_nodes(h_nodes, h_neighbors, E_idx):
    h_nodes = gather_nodes(h_nodes, E_idx)
    h_nn = torch.cat([h_neighbors, h_nodes], -1)
    return h_nn


class Normalize(nn.Module): 
    def __init__(self, features, epsilon=1e-6):
        super(Normalize, self).__init__()
        self.gain = nn.Parameter(torch.ones(features))
        self.bias = nn.Parameter(torch.zeros(features))
        self.epsilon = epsilon

    def forward(self, x, dim=-1):
        mu = x.mean(dim, keepdim=True)
        sigma = torch.sqrt(x.var(dim, keepdim=True) + self.epsilon)
        gain = self.gain
        bias = self.bias
        # Reshape
        if dim != -1:
            shape = [1] * len(mu.size())
            shape[dim] = self.gain.size()[0]
            gain = gain.view(shape)
            bias = bias.view(shape)
        return gain * (x - mu) / (sigma + self.epsilon) + bias


class TransformerLayer(nn.Module):
    def __init__(self, num_hidden, num_in, num_heads=4, dropout=0.2):
        super(TransformerLayer, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.relu = nn.LeakyReLU()
        self.attn_norm = nn.LayerNorm(num_hidden, eps=1e-6)
        self.norm = nn.ModuleList([Normalize(num_hidden) for _ in range(3)])
        self.ffn_norm = nn.LayerNorm(num_hidden, eps=1e-6)

        self.attention = NeighborAttention(num_hidden, num_in, num_heads)
        self.dense1 = PositionWiseFeedForward(num_hidden, num_hidden * 4)
        self.dense2 = PositionWiseFeedForward(num_hidden, num_hidden * 4)
        self.dense3 = PositionWiseFeedForward(num_hidden, num_hidden * 4)
        self.weight1 = Parameter(torch.FloatTensor(128 * 2, 128))
        self.weight2 = Parameter(torch.FloatTensor(128 * 2, 128))
        self.gcn1 = GCN(128)
        self.gcn2 = GCN(128)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(128)
        self.weight1.data.uniform_(-stdv, stdv)
        self.weight2.data.uniform_(-stdv, stdv)

    def forward(self, h_V, h_E,adj,h0,l,mask_V=None, mask_attend=None):    # mask_attend [B, L, K]
        # Concatenate h_V_i to h_E_ij      
        """ Parallel computation of full transformer layer """
        # Self-attention
        # h_V = self.attn_norm(h_V)
        # h_V = self.relu(h_V)
        first_dimension = h_V.shape[0]
        theta = min(1, math.log(1.5 / l + 1))

        dh = self.attention(h_V, h_E, mask_attend)
        h_V = self.norm[0](h_V + self.dropout(dh))
        h_V = self.relu(h_V)

        # # theta = min(1, math.log(1.5 / l + 1))
        dh = self.gcn1(h_V, adj, mask_V)
        support = torch.cat([dh, h0], -1)
        # support = support.view(-1, 64)
        r = (1 - 0.7) * dh + 0.7 * h0
        weight_expanded = self.weight1.unsqueeze(0).expand(first_dimension, 256, 128)
        output = theta * torch.bmm(support, weight_expanded) + (1 - theta) * r
        h_V = self.norm[1](h_V + self.dropout(self.relu(output)))


        dh = self.dense1(h_V)
        h_V = self.norm[2](h_V + self.dropout(dh))


        if mask_V is not None: 
            mask_V = mask_V.unsqueeze(-1)
            h_V = mask_V * h_V
        return h_V


class PositionWiseFeedForward(nn.Module):
    def __init__(self, num_hidden, num_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.W_in = nn.Linear(num_hidden, num_ff, bias=True)
        self.W_out = nn.Linear(num_ff, num_hidden, bias=True)

    def forward(self, h_V):
        h = F.relu(self.W_in(h_V))
        h = self.W_out(h)
        return h


class NeighborAttention(nn.Module):
    def __init__(self, num_hidden, num_in, num_heads=4):
        super(NeighborAttention, self).__init__()
        self.num_heads = num_heads
        self.num_hidden = num_hidden

        # Self-attention layers: {queries, keys, values, output}
        self.W_Q = nn.Linear(num_hidden, num_hidden, bias=False)
        self.W_K = nn.Linear(num_in, num_hidden, bias=False)
        self.W_V = nn.Linear(num_in, num_hidden, bias=False)
        self.W_O = nn.Linear(num_hidden, num_hidden, bias=False)

    def _masked_softmax(self, attend_logits, mask_attend, dim=-1):
        """ Numerically stable masked softmax """
        negative_inf = np.finfo(np.float32).min
        attend_logits = torch.where(
            mask_attend > 0,
            attend_logits,
            attend_logits.new_tensor(negative_inf),
        )
        attend = F.softmax(attend_logits, dim)
        attend = mask_attend * attend
        return attend

    def forward(self, h_V, h_E, mask_attend=None):
        """ Self-attention, graph-structured O(Nk)
        Args:
            h_V:            Node features           [N_batch, N_nodes, N_hidden]
            h_E:            Neighbor features       [N_batch, N_nodes, K, N_hidden] 
            mask_attend:    Mask for attention      [N_batch, N_nodes, K]           mask_attend [B, L, K] 
        Returns:
            h_V:            Node update
        """

        # Queries, Keys, Values
        n_batch, n_nodes, n_neighbors = h_E.shape[:3]
        n_heads = self.num_heads

        d = int(self.num_hidden / n_heads)
        Q = self.W_Q(h_V).view([n_batch, n_nodes, 1, n_heads, 1, d])
        K = self.W_K(h_E).view([n_batch, n_nodes, n_neighbors, n_heads, d, 1])
        V = self.W_V(h_E).view([n_batch, n_nodes, n_neighbors, n_heads, d])

        # Attention with scaled inner product
        attend_logits = torch.matmul(Q, K).view([n_batch, n_nodes, n_neighbors, n_heads]).transpose(-2,-1)
        attend_logits = attend_logits / np.sqrt(d)
        
        if mask_attend is not None:
            # Masked softmax
            mask = mask_attend.unsqueeze(2).expand(-1,-1,n_heads,-1)
            attend = self._masked_softmax(attend_logits, mask) # [B, L, heads, K]
        else:
            attend = F.softmax(attend_logits, -1)

        # Attentive reduction
        h_V_update = torch.matmul(attend.unsqueeze(-2), V.transpose(2,3)) # [B, L, heads, 1, K] × [B, L, heads, K, d]
        h_V_update = h_V_update.view([n_batch, n_nodes, self.num_hidden])
        h_V_update = self.W_O(h_V_update)
        return h_V_update
class GCN(nn.Module):
    def __init__(self, hidden_size):
        super(GCN, self).__init__()
        self.hidden_size = hidden_size
        self.weight = Parameter(torch.FloatTensor(hidden_size, hidden_size))
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.hidden_size)
        self.weight.data.uniform_(-stdv, stdv)

    def forward(self, x, adj, mask):
        # mask shape: (B, N), where 1 indicates valid nodes and 0 indicates padded nodes
        B, N, D = x.size()

        # Expand weight for batch multiplication
        expanded_weight = self.weight.unsqueeze(0).expand(B, -1, -1)

        # Apply mask to x and adj
        x = x * mask.unsqueeze(-1)
        adj = adj * mask.unsqueeze(-1) * mask.unsqueeze(1)

        # Perform graph convolution
        y = torch.bmm(adj, x)
        y = torch.bmm(y, expanded_weight)

        # Apply mask to output
        y = y * mask.unsqueeze(-1)

        return y

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torch.nn.parameter import Parameter
	import math
	# The following gather functions
	def gather_edges(edges, neighbor_idx):
	# Features [B,N,N,C] at Neighbor indices [B,N,K] => Neighbor features [B,N,K,C]
	neighbors = neighbor_idx.unsqueeze(-1).expand(-1, -1, -1, edges.size(-1))
	edge_features = torch.gather(edges, 2, neighbors)
	return edge_features

	def gather_nodes(nodes, neighbor_idx):
	# Features [B,N,C] at Neighbor indices [B,N,K] => [B,N,K,C]
	# Flatten and expand indices per batch [B,N,K] => [B,NK] => [B,NK,C]
	neighbors_flat = neighbor_idx.view((neighbor_idx.shape[0], -1))
	neighbors_flat = neighbors_flat.unsqueeze(-1).expand(-1, -1, nodes.size(2))
	# Gather and re-pack
	neighbor_features = torch.gather(nodes, 1, neighbors_flat)
	neighbor_features = neighbor_features.view(list(neighbor_idx.shape)[:3] + [-1])
	return neighbor_features

	def cat_neighbors_nodes(h_nodes, h_neighbors, E_idx):
	h_nodes = gather_nodes(h_nodes, E_idx)
	h_nn = torch.cat([h_neighbors, h_nodes], -1)
	return h_nn


	class Normalize(nn.Module):
	def __init__(self, features, epsilon=1e-6):
	super(Normalize, self).__init__()
	self.gain = nn.Parameter(torch.ones(features))
	self.bias = nn.Parameter(torch.zeros(features))
	self.epsilon = epsilon

	def forward(self, x, dim=-1):
	mu = x.mean(dim, keepdim=True)
	sigma = torch.sqrt(x.var(dim, keepdim=True) + self.epsilon)
	gain = self.gain
	bias = self.bias
	# Reshape
	if dim != -1:
	shape = [1] * len(mu.size())
	shape[dim] = self.gain.size()[0]
	gain = gain.view(shape)
	bias = bias.view(shape)
	return gain * (x - mu) / (sigma + self.epsilon) + bias


	class TransformerLayer(nn.Module):
	def __init__(self, num_hidden, num_in, num_heads=4, dropout=0.2):
	super(TransformerLayer, self).__init__()
	self.dropout = nn.Dropout(dropout)
	self.relu = nn.LeakyReLU()
	self.attn_norm = nn.LayerNorm(num_hidden, eps=1e-6)
	self.norm = nn.ModuleList([Normalize(num_hidden) for _ in range(3)])
	self.ffn_norm = nn.LayerNorm(num_hidden, eps=1e-6)

	self.attention = NeighborAttention(num_hidden, num_in, num_heads)
	self.dense1 = PositionWiseFeedForward(num_hidden, num_hidden * 4)
	self.dense2 = PositionWiseFeedForward(num_hidden, num_hidden * 4)
	self.dense3 = PositionWiseFeedForward(num_hidden, num_hidden * 4)
	self.weight1 = Parameter(torch.FloatTensor(128 * 2, 128))
	self.weight2 = Parameter(torch.FloatTensor(128 * 2, 128))
	self.gcn1 = GCN(128)
	self.gcn2 = GCN(128)
	self.reset_parameters()

	def reset_parameters(self):
	stdv = 1. / math.sqrt(128)
	self.weight1.data.uniform_(-stdv, stdv)
	self.weight2.data.uniform_(-stdv, stdv)

	def forward(self, h_V, h_E,adj,h0,l,mask_V=None, mask_attend=None): # mask_attend [B, L, K]
	# Concatenate h_V_i to h_E_ij
	""" Parallel computation of full transformer layer """
	# Self-attention
	# h_V = self.attn_norm(h_V)
	# h_V = self.relu(h_V)
	first_dimension = h_V.shape[0]
	theta = min(1, math.log(1.5 / l + 1))

	dh = self.attention(h_V, h_E, mask_attend)
	h_V = self.norm[0](h_V + self.dropout(dh))
	h_V = self.relu(h_V)

	# # theta = min(1, math.log(1.5 / l + 1))
	dh = self.gcn1(h_V, adj, mask_V)
	support = torch.cat([dh, h0], -1)
	# support = support.view(-1, 64)
	r = (1 - 0.7) * dh + 0.7 * h0
	weight_expanded = self.weight1.unsqueeze(0).expand(first_dimension, 256, 128)
	output = theta * torch.bmm(support, weight_expanded) + (1 - theta) * r
	h_V = self.norm[1](h_V + self.dropout(self.relu(output)))


	dh = self.dense1(h_V)
	h_V = self.norm[2](h_V + self.dropout(dh))


	if mask_V is not None:
	mask_V = mask_V.unsqueeze(-1)
	h_V = mask_V * h_V
	return h_V


	class PositionWiseFeedForward(nn.Module):
	def __init__(self, num_hidden, num_ff):
	super(PositionWiseFeedForward, self).__init__()
	self.W_in = nn.Linear(num_hidden, num_ff, bias=True)
	self.W_out = nn.Linear(num_ff, num_hidden, bias=True)

	def forward(self, h_V):
	h = F.relu(self.W_in(h_V))
	h = self.W_out(h)
	return h


	class NeighborAttention(nn.Module):
	def __init__(self, num_hidden, num_in, num_heads=4):
	super(NeighborAttention, self).__init__()
	self.num_heads = num_heads
	self.num_hidden = num_hidden

	# Self-attention layers: {queries, keys, values, output}
	self.W_Q = nn.Linear(num_hidden, num_hidden, bias=False)
	self.W_K = nn.Linear(num_in, num_hidden, bias=False)
	self.W_V = nn.Linear(num_in, num_hidden, bias=False)
	self.W_O = nn.Linear(num_hidden, num_hidden, bias=False)

	def _masked_softmax(self, attend_logits, mask_attend, dim=-1):
	""" Numerically stable masked softmax """
	negative_inf = np.finfo(np.float32).min
	attend_logits = torch.where(
	mask_attend > 0,
	attend_logits,
	attend_logits.new_tensor(negative_inf),
	)
	attend = F.softmax(attend_logits, dim)
	attend = mask_attend * attend
	return attend

	def forward(self, h_V, h_E, mask_attend=None):
	""" Self-attention, graph-structured O(Nk)
	Args:
	h_V: Node features [N_batch, N_nodes, N_hidden]
	h_E: Neighbor features [N_batch, N_nodes, K, N_hidden]
	mask_attend: Mask for attention [N_batch, N_nodes, K] mask_attend [B, L, K]
	Returns:
	h_V: Node update
	"""

	# Queries, Keys, Values
	n_batch, n_nodes, n_neighbors = h_E.shape[:3]
	n_heads = self.num_heads

	d = int(self.num_hidden / n_heads)
	Q = self.W_Q(h_V).view([n_batch, n_nodes, 1, n_heads, 1, d])
	K = self.W_K(h_E).view([n_batch, n_nodes, n_neighbors, n_heads, d, 1])
	V = self.W_V(h_E).view([n_batch, n_nodes, n_neighbors, n_heads, d])

	# Attention with scaled inner product
	attend_logits = torch.matmul(Q, K).view([n_batch, n_nodes, n_neighbors, n_heads]).transpose(-2,-1)
	attend_logits = attend_logits / np.sqrt(d)

	if mask_attend is not None:
	# Masked softmax
	mask = mask_attend.unsqueeze(2).expand(-1,-1,n_heads,-1)
	attend = self._masked_softmax(attend_logits, mask) # [B, L, heads, K]
	else:
	attend = F.softmax(attend_logits, -1)

	# Attentive reduction
	h_V_update = torch.matmul(attend.unsqueeze(-2), V.transpose(2,3)) # [B, L, heads, 1, K] × [B, L, heads, K, d]
	h_V_update = h_V_update.view([n_batch, n_nodes, self.num_hidden])
	h_V_update = self.W_O(h_V_update)
	return h_V_update
	class GCN(nn.Module):
	def __init__(self, hidden_size):
	super(GCN, self).__init__()
	self.hidden_size = hidden_size
	self.weight = Parameter(torch.FloatTensor(hidden_size, hidden_size))
	self.reset_parameters()

	def reset_parameters(self):
	stdv = 1. / math.sqrt(self.hidden_size)
	self.weight.data.uniform_(-stdv, stdv)

	def forward(self, x, adj, mask):
	# mask shape: (B, N), where 1 indicates valid nodes and 0 indicates padded nodes
	B, N, D = x.size()

	# Expand weight for batch multiplication
	expanded_weight = self.weight.unsqueeze(0).expand(B, -1, -1)

	# Apply mask to x and adj
	x = x * mask.unsqueeze(-1)
	adj = adj * mask.unsqueeze(-1) * mask.unsqueeze(1)

	# Perform graph convolution
	y = torch.bmm(adj, x)
	y = torch.bmm(y, expanded_weight)

	# Apply mask to output
	y = y * mask.unsqueeze(-1)

	return y