model_hyper.py

import torch 
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.nn.utils.rnn import pad_sequence
from torch_geometric.nn import RGCNConv, GraphConv, FAConv
from torch.nn import Parameter
import numpy as np, itertools, random, copy, math
import math
import scipy.sparse as sp
from model_GCN import GCNII_lyc
import ipdb
from HypergraphConv import HypergraphConv
from torch_geometric.nn import GCNConv
from itertools import permutations
from torch_geometric.nn.pool.topk_pool import topk
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp, global_add_pool as gsp
from torch_geometric.nn.inits import glorot
from torch_geometric.nn.conv.gcn_conv import gcn_norm
from torch_geometric.utils import add_self_loops
from high_fre_conv import highConv
from torch_geometric.nn import TransformerConv
class STEFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        return (input > 0).float()

    @staticmethod
    def backward(ctx, grad_output):
        return F.hardtanh(grad_output)

class GraphConvolution(nn.Module):

    def __init__(self, in_features, out_features, residual=False, variant=False):
        super(GraphConvolution, self).__init__() 
        self.variant = variant
        if self.variant:
            self.in_features = 2*in_features 
        else:
            self.in_features = in_features

        self.out_features = out_features
        self.residual = residual
        self.weight = Parameter(torch.FloatTensor(self.in_features,self.out_features))
        self.reset_parameters()

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

    def forward(self, input, adj , h0 , lamda, alpha, l):
        theta = math.log(lamda/l+1)
        hi = torch.spmm(adj, input)
        if self.variant:
            support = torch.cat([hi,h0],1)
            r = (1-alpha)*hi+alpha*h0
        else:
            support = (1-alpha)*hi+alpha*h0
            r = support
        output = theta*torch.mm(support, self.weight)+(1-theta)*r
        if self.residual:
            output = output+input
        return output
class SGCN(nn.Module):

    def __init__(self, g_dim, h2_dim):
        super(SGCN, self).__init__()
        self.conv1 = TransformerConv(g_dim, h2_dim)


    def forward(self, node_features, edge_index):
        x = self.conv1(node_features, edge_index)
        return x
class emoAgent(nn.Module):
    def __init__(self, a_dim, v_dim, l_dim, n_dim, nlayers, nhidden, nclass, dropout, lamda, alpha, variant, return_feature, use_residue, 
                new_graph='full',n_speakers=2, modals=['a','v','l'], use_speaker=True, use_modal=False, num_L=3, num_K=4):
        super(emoAgent, self).__init__()
        self.return_feature = return_feature  #True
        self.use_residue = use_residue
        self.new_graph = new_graph
        self.nhidden = nhidden
        self.act_fn = nn.ReLU()
        self.dropout = dropout
        self.alpha = alpha
        self.lamda = lamda
        self.modals = modals
        self.modal_embeddings = nn.Embedding(3, n_dim)
        self.speaker_embeddings = nn.Embedding(n_speakers, n_dim)
        self.use_speaker = use_speaker
        self.use_modal = use_modal
        self.use_position = False
        #------------------------------------    
        self.fc1 = nn.Linear(n_dim, nhidden)  
        self.fc2 = nn.Linear(nhidden * 3, nhidden)  
        self.num_L =  num_L
        self.num_K =  num_K
        for ll in range(num_L):
            setattr(self,'hyperconv%d' %(ll+1), HypergraphConv(nhidden, nhidden))
        self.act_fn = nn.ReLU()
        self.hyperedge_weight = nn.Parameter(torch.ones(10000))
        self.EW_weight = nn.Parameter(torch.ones(20000))
        self.hyperedge_attr1 = nn.Parameter(torch.rand(nhidden))
        self.hyperedge_attr2 = nn.Parameter(torch.rand(nhidden))
        self.hyperedge_attr3 = nn.Parameter(torch.rand(nhidden))
        

    def forward(self, a, v, l, agent_gnn, dia_len, qmask, epoch):
        hyperedge_index, edge_index, features, batch, hyperedge_type1 = self.create_hyper_index(a, v, l, agent_gnn, dia_len, self.modals)
        x1 = self.fc1(features)  
        weight = self.hyperedge_weight[0:hyperedge_index[1].max().item()+1]
        EW_weight = self.EW_weight[0:hyperedge_index.size(1)]
        edge_attr = self.hyperedge_attr1*hyperedge_type1 + self.hyperedge_attr2*((1+hyperedge_type1)//3) + self.hyperedge_attr3*(2-hyperedge_type1)
        out = x1
        for ll in range(self.num_L): # num_L
            out = getattr(self,'hyperconv%d' %(ll+1))(out, hyperedge_index, weight, edge_attr, EW_weight, dia_len)  
        out = out + x1
        
        if self.use_residue:
            out = torch.cat([features, out], dim=-1)    
        out = self.reverse_features(dia_len, out)
        return out 
        
    def create_hyper_index(self, a, v, l, agent, dia_len, modals):
        self_loop = False
        num_modality = len(modals) + 1
        node_count = 0
        edge_count = 0
        batch_count = 0
        index1 =[]
        index2 =[]
        tmp = []
        batch = []
        hyperedge_type1 = []
        for i in dia_len:
            nodes = list(range(i*num_modality)) 
            nodes = [j + node_count for j in nodes]
            nodes_l = nodes[0:i*num_modality//4]
            nodes_a = nodes[i*num_modality//4:i*num_modality*2//4]
            nodes_v = nodes[i*num_modality*2//4:i*num_modality*3//4]
            nodes_agent = nodes[i*num_modality*3//4:]
            index1 = index1 + nodes_l + nodes_a + nodes_v + nodes_agent
            for _ in range(i):
                index1 = index1 + [nodes_l[_]] + [nodes_a[_]] + [nodes_agent[_]] 
                index1 = index1 + [nodes_l[_]] + [nodes_v[_]] + [nodes_agent[_]] 
                index1 = index1 + [nodes_a[_]] + [nodes_v[_]] + [nodes_agent[_]] 
                index1 = index1 + [nodes_a[_]] + [nodes_agent[_]] 
                index1 = index1 + [nodes_v[_]] + [nodes_agent[_]] 
                index1 = index1 + [nodes_l[_]] + [nodes_agent[_]] 
            for _ in range(i+4):
                if _ < 4:
                    index2 = index2 + [edge_count]*i
                else:                 
                    index2 = index2 + [edge_count]*3
                    edge_count = edge_count + 1
                    index2 = index2 + [edge_count]*3
                    edge_count = edge_count + 1
                    index2 = index2 + [edge_count]*3
                    edge_count = edge_count + 1
                    index2 = index2 + [edge_count]*2
                    edge_count = edge_count + 1
                    index2 = index2 + [edge_count]*2   
                    edge_count = edge_count + 1
                    index2 = index2 + [edge_count]*2                   
                edge_count = edge_count + 1 
            if node_count == 0:
                ll = l[0:0+i]
                aa = a[0:0+i]
                vv = v[0:0+i]
                gg = agent[0:0+i]
                features = torch.cat([ll,aa,vv,gg],dim=0)
                temp = 0+i
            else:
                ll = l[temp:temp+i]
                aa = a[temp:temp+i]
                vv = v[temp:temp+i]
                gg = agent[temp:temp+i]
                features_temp = torch.cat([ll,aa,vv,gg],dim=0)
                features =  torch.cat([features,features_temp],dim=0)
                temp = temp + i
            
            Gnodes=[]
            Gnodes.append(nodes_l)
            Gnodes.append(nodes_a)
            Gnodes.append(nodes_v)
            Gnodes.append(nodes_agent)
            for _ in range(i):
                Gnodes.append([nodes_l[_]] + [nodes_agent[_]] )
            for _ in range(i):
                Gnodes.append([nodes_a[_]] + [nodes_agent[_]] )
            for _ in range(i):
                Gnodes.append([nodes_v[_]] + [nodes_agent[_]] )
            for _ in range(i):
                Gnodes.append([nodes_l[_]] + [nodes_a[_]] + [nodes_agent[_]] )
            for _ in range(i):
                Gnodes.append([nodes_l[_]] + [nodes_v[_]] + [nodes_agent[_]] )
            for _ in range(i):
                Gnodes.append([nodes_a[_]] + [nodes_v[_]] + [nodes_agent[_]] )
            hyperedge_type1 = hyperedge_type1 + [2]*3*i + [1]*3*i + [0]*4

            node_count = node_count + i*num_modality

        index1 = torch.LongTensor(index1).view(1,-1)
        index2 = torch.LongTensor(index2).view(1,-1)
        hyperedge_index = torch.cat([index1,index2],dim=0).cuda() 
        if self_loop:
            max_edge = hyperedge_index[1].max()
            max_node = hyperedge_index[0].max()
            loops = torch.cat([torch.arange(0,max_node+1,1).repeat_interleave(2).view(1,-1),
                            torch.arange(max_edge+1,max_edge+1+max_node+1,1).repeat_interleave(2).view(1,-1)],dim=0).cuda()
            hyperedge_index = torch.cat([hyperedge_index, loops], dim=1)

        edge_index = torch.LongTensor(tmp).T.cuda()

        hyperedge_type1 = torch.LongTensor(hyperedge_type1).view(-1,1).cuda()

        return hyperedge_index, edge_index, features, batch, hyperedge_type1

    
    def reverse_features(self, dia_len, features):
        l=[]
        a=[]
        v=[]
        g=[]
        for i in dia_len:
            ll = features[0:1*i]
            aa = features[1*i:2*i]
            vv = features[2*i:3*i]
            gg = features[3*i:4*i]
            features = features[4*i:]
            l.append(ll)
            a.append(aa)
            v.append(vv)
            g.append(gg)
        tmpl = torch.cat(l,dim=0)
        tmpa = torch.cat(a,dim=0)
        tmpv = torch.cat(v,dim=0)
        tmpg = torch.cat(g,dim=0)
        features = torch.cat([tmpl, tmpa, tmpv, tmpg], dim=-1)
        return features

class emoGen(nn.Module):
    def __init__(self, a_dim, v_dim, l_dim, n_dim, nlayers, nhidden, nclass, dropout, lamda, alpha, variant, return_feature, use_residue, 
                new_graph='full',n_speakers=2, modals=['a','v','l'], use_speaker=True, use_modal=False, num_L=3, num_K=4):
        super(emoGen, self).__init__()
        self.return_feature = return_feature  #True
        self.use_residue = use_residue
        self.new_graph = new_graph
        self.nhidden = nhidden
        self.act_fn = nn.ReLU()
        self.dropout = dropout
        self.alpha = alpha
        self.lamda = lamda
        self.modals = modals
        self.modal_embeddings = nn.Embedding(3, n_dim)
        self.speaker_embeddings = nn.Embedding(n_speakers, n_dim)
        self.use_speaker = use_speaker
        self.use_modal = use_modal
        self.use_position = False
        self.fc1 = nn.Linear(n_dim, nhidden)  
        self.fc2 = nn.Linear(nhidden * 3, nhidden)    
        self.num_L =  num_L
        self.num_K =  num_K
        for kk in range(num_K):
            setattr(self,'conv%d' %(kk+1), highConv(nhidden, nhidden))
        

    def forward(self, a, v, l, dia_len, qmask, epoch):
        qmask = torch.cat([qmask[:x,i,:] for i,x in enumerate(dia_len)],dim=0)
        spk_idx = torch.argmax(qmask, dim=-1)
        spk_emb_vector = self.speaker_embeddings(spk_idx)
        if self.use_speaker:
            if 'l' in self.modals:
                l += spk_emb_vector
        if self.use_position:
            if 'l' in self.modals:
                l = self.l_pos(l, dia_len)
            if 'a' in self.modals:
                a = self.a_pos(a, dia_len)
            if 'v' in self.modals:
                v = self.v_pos(v, dia_len)
        if self.use_modal:  
            emb_idx = torch.LongTensor([0, 1, 2]).cuda()
            emb_vector = self.modal_embeddings(emb_idx)
            if 'a' in self.modals:
                a += emb_vector[0].reshape(1, -1).expand(a.shape[0], a.shape[1])
            if 'v' in self.modals:
                v += emb_vector[1].reshape(1, -1).expand(v.shape[0], v.shape[1])
            if 'l' in self.modals:
                l += emb_vector[2].reshape(1, -1).expand(l.shape[0], l.shape[1])
        gnn_edge_index, gnn_features = self.create_gnn_index(a, v, l, dia_len, self.modals)
        x1 = self.fc1(gnn_features)  
        out = x1
        gnn_out = x1
        for kk in range(self.num_K):
            gnn_out = gnn_out + getattr(self,'conv%d' %(kk+1))(gnn_out,gnn_edge_index)
        out = out + gnn_out
        if self.use_residue:
            out = torch.cat([gnn_features, out], dim=-1)
        agent = self.reverse_features(dia_len, out) 
        agent = self.fc2(agent)
        return agent 

    def reverse_features(self, dia_len, features):
        l=[]
        a=[]
        v=[]
        for i in dia_len:
            ll = features[0:1*i]
            aa = features[1*i:2*i]
            vv = features[2*i:3*i]
            features = features[3*i:]
            l.append(ll)
            a.append(aa)
            v.append(vv)
        tmpl = torch.cat(l,dim=0)
        tmpa = torch.cat(a,dim=0)
        tmpv = torch.cat(v,dim=0)
        features = torch.cat([tmpl, tmpa, tmpv], dim=-1)
        return features

    def create_gnn_index(self, a, v, l, dia_len, modals):
        self_loop = False
        num_modality = len(modals)
        node_count = 0
        batch_count = 0
        index =[]
        tmp = []
        
        for i in dia_len:
            nodes = list(range(i*num_modality))
            nodes = [j + node_count for j in nodes]
            nodes_l = nodes[0:i*num_modality//3]
            nodes_a = nodes[i*num_modality//3:i*num_modality*2//3]
            nodes_v = nodes[i*num_modality*2//3:]
            index = index + list(permutations(nodes_l,2)) + list(permutations(nodes_a,2)) + list(permutations(nodes_v,2))
            Gnodes=[]
            for _ in range(i):
                Gnodes.append([nodes_l[_]] + [nodes_a[_]] + [nodes_v[_]])
            for ii, _ in enumerate(Gnodes):
                tmp = tmp +  list(permutations(_,2))
            if node_count == 0:
                ll = l[0:0+i]
                aa = a[0:0+i]
                vv = v[0:0+i]
                features = torch.cat([ll,aa,vv],dim=0)
                temp = 0+i
            else:
                ll = l[temp:temp+i]
                aa = a[temp:temp+i]
                vv = v[temp:temp+i]
                features_temp = torch.cat([ll,aa,vv],dim=0)
                features =  torch.cat([features,features_temp],dim=0)
                temp = temp+i
            node_count = node_count + i*num_modality
        edge_index = torch.cat([torch.LongTensor(index).T,torch.LongTensor(tmp).T],1).cuda()

        return edge_index, features