[Example] TAHIN (#2864)

* tahin

* readme

* readme

* readme

* readme

* readme

* readme

* main

* main

* new_line

* update
Co-authored-by: zhjwy9343 <6593865@qq.com>

examples/README.md

@@ -93,6 +93,7 @@ The folder contains example implementations of selected research papers related
 | [GNNExplainer: Generating Explanations for Graph Neural Networks](#gnnexplainer) |  :heavy_check_mark: |                                  |                                  |                                  |                                  |
 | [Interaction Networks for Learning about Objects, Relations and Physics](#graphsim) |  |                                 |:heavy_check_mark:                 |                    |                    |
 | [Representation Learning on Graphs with Jumping Knowledge Networks](#jknet) |  :heavy_check_mark: |                                  |                                  |                                  |                                  |
+| [A Heterogeneous Information Network based Cross Domain Insurance Recommendation System for Cold Start Users](#tahin) |  |     :heavy_check_mark:       |                 |                    |                    |
 | [DeeperGCN: All You Need to Train Deeper GCNs](#deepergcn)                                                                        |                                  |                                  | :heavy_check_mark:        |                    | :heavy_check_mark: |
 | [Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forcasting](#dcrnn) |                     |                                  | :heavy_check_mark:        |                    |                    |
 | [GaAN: Gated Attention Networks for Learning on large and Spatiotemporal Graphs](#gaan) |                     |                                  | :heavy_check_mark:        |                    |                    |
@@ -156,6 +157,10 @@ The folder contains example implementations of selected research papers related
     - Example code: [PyTorch](../examples/pytorch/deepergcn)
     - Tags: over-smoothing, deeper gnn, OGB
 
+- <a name="tahin"></a> Bi, Ye, et al. A Heterogeneous Information Network based Cross DomainInsurance Recommendation System for Cold Start Users. [Paper link](https://arxiv.org/abs/2007.15293).
+    - Example code: [Pytorch](../examples/pytorch/TAHIN)
+    - Tags: cross-domain recommendation, graph neural network
+
 ## 2019
 
 - <a name="infograph"></a> Sun et al. InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization. [Paper link](https://arxiv.org/abs/1908.01000). 

examples/pytorch/TAHIN/TAHIN.py

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import dgl
+import dgl.function as fn
+from dgl.nn.pytorch import GATConv
+
+#Semantic attention in the metapath-based aggregation (the same as that in the HAN)
+class SemanticAttention(nn.Module):
+    def __init__(self, in_size, hidden_size=128):
+        super(SemanticAttention, self).__init__()
+
+        self.project = nn.Sequential(
+            nn.Linear(in_size, hidden_size),
+            nn.Tanh(),
+            nn.Linear(hidden_size, 1, bias=False)
+        )
+
+    def forward(self, z):
+        '''
+        Shape of z: (N, M , D*K)
+        N: number of nodes
+        M: number of metapath patterns
+        D: hidden_size
+        K: number of heads
+        '''
+        w = self.project(z).mean(0)                    # (M, 1)
+        beta = torch.softmax(w, dim=0)                 # (M, 1)
+        beta = beta.expand((z.shape[0],) + beta.shape) # (N, M, 1)
+
+        return (beta * z).sum(1)                       # (N, D * K)
+
+#Metapath-based aggregation (the same as the HANLayer)
+class HANLayer(nn.Module):
+    def __init__(self, meta_path_patterns, in_size, out_size, layer_num_heads, dropout):
+        super(HANLayer, self).__init__()
+
+        # One GAT layer for each meta path based adjacency matrix
+        self.gat_layers = nn.ModuleList()
+        for i in range(len(meta_path_patterns)):
+            self.gat_layers.append(GATConv(in_size, out_size, layer_num_heads,
+                                           dropout, dropout, activation=F.elu,
+                                           allow_zero_in_degree=True))
+        self.semantic_attention = SemanticAttention(in_size=out_size * layer_num_heads)
+        self.meta_path_patterns = list(tuple(meta_path_pattern) for meta_path_pattern in meta_path_patterns)
+
+        self._cached_graph = None
+        self._cached_coalesced_graph = {}
+
+    def forward(self, g, h):
+        semantic_embeddings = []
+        #obtain metapath reachable graph
+        if self._cached_graph is None or self._cached_graph is not g:
+            self._cached_graph = g
+            self._cached_coalesced_graph.clear()
+            for meta_path_pattern in self.meta_path_patterns:
+                self._cached_coalesced_graph[meta_path_pattern] = dgl.metapath_reachable_graph(
+                        g, meta_path_pattern)
+
+        for i, meta_path_pattern in enumerate(self.meta_path_patterns):
+            new_g = self._cached_coalesced_graph[meta_path_pattern]
+            semantic_embeddings.append(self.gat_layers[i](new_g, h).flatten(1))
+        semantic_embeddings = torch.stack(semantic_embeddings, dim=1)                  # (N, M, D * K)
+
+        return self.semantic_attention(semantic_embeddings)                            # (N, D * K)
+
+#Relational neighbor aggregation
+class RelationalAGG(nn.Module):
+    def __init__(self, g, in_size, out_size, dropout=0.1):
+        super(RelationalAGG, self).__init__()
+        self.in_size = in_size
+        self.out_size = out_size
+
+        #Transform weights for different types of edges
+        self.W_T = nn.ModuleDict({
+            name : nn.Linear(in_size, out_size, bias = False) for name in g.etypes
+        })
+
+        #Attention weights for different types of edges
+        self.W_A = nn.ModuleDict({
+            name : nn.Linear(out_size, 1, bias = False) for name in g.etypes
+        })
+
+        #layernorm
+        self.layernorm = nn.LayerNorm(out_size)
+
+        #dropout layer
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, g, feat_dict):
+        funcs={}
+        for srctype, etype, dsttype in g.canonical_etypes:
+            g.nodes[dsttype].data['h'] = feat_dict[dsttype] #nodes' original feature
+            g.nodes[srctype].data['h'] = feat_dict[srctype]  
+            g.nodes[srctype].data['t_h'] = self.W_T[etype](feat_dict[srctype])  #src nodes' transformed feature
+
+            #compute the attention numerator (exp)
+            g.apply_edges(fn.u_mul_v('t_h','h','x'),etype=etype)
+            g.edges[etype].data['x'] = torch.exp(self.W_A[etype](g.edges[etype].data['x']))
+
+            #first update to compute the attention denominator (\sum exp)
+            funcs[etype] = (fn.copy_e('x', 'm'), fn.sum('m', 'att'))  
+        g.multi_update_all(funcs, 'sum') 
+
+        funcs={}
+        for srctype, etype, dsttype in g.canonical_etypes:
+            g.apply_edges(fn.e_div_v('x', 'att', 'att'),etype=etype) #compute attention weights (numerator/denominator)
+            funcs[etype] = (fn.u_mul_e('h', 'att', 'm'), fn.sum('m', 'h'))  #\sum(h0*att) -> h1
+        #second update to obtain h1
+        g.multi_update_all(funcs, 'sum') 
+
+        #apply activation, layernorm, and dropout
+        feat_dict={}
+        for ntype in g.ntypes:
+            feat_dict[ntype] = self.dropout(self.layernorm(F.relu_(g.nodes[ntype].data['h']))) #apply activation, layernorm, and dropout
+        
+        return feat_dict
+
+class TAHIN(nn.Module):
+    def __init__(self, g, meta_path_patterns, in_size, out_size, num_heads, dropout):
+        super(TAHIN, self).__init__()
+
+        #embeddings for different types of nodes, h0
+        self.initializer = nn.init.xavier_uniform_
+        self.feature_dict = nn.ParameterDict({
+            ntype: nn.Parameter(self.initializer(torch.empty(g.num_nodes(ntype), in_size))) for ntype in g.ntypes
+        })
+
+        #relational neighbor aggregation, this produces h1
+        self.RelationalAGG = RelationalAGG(g, in_size, out_size)
+
+        #metapath-based aggregation modules for user and item, this produces h2
+        self.meta_path_patterns = meta_path_patterns 
+        #one HANLayer for user, one HANLayer for item
+        self.hans = nn.ModuleDict({
+            key: HANLayer(value, in_size, out_size, num_heads, dropout) for key, value in self.meta_path_patterns.items()
+        })
+
+        #layers to combine h0, h1, and h2
+        #used to update node embeddings
+        self.user_layer1 = nn.Linear((num_heads+1)*out_size, out_size, bias=True)
+        self.user_layer2 = nn.Linear(2*out_size, out_size, bias=True)
+        self.item_layer1 = nn.Linear((num_heads+1)*out_size, out_size, bias=True)
+        self.item_layer2 = nn.Linear(2*out_size, out_size, bias=True)
+
+        #layernorm
+        self.layernorm = nn.LayerNorm(out_size)
+
+        #network to score the node pairs
+        self.pred = nn.Linear(out_size, out_size)
+        self.dropout = nn.Dropout(dropout)
+        self.fc = nn.Linear(out_size, 1)
+
+    def forward(self, g, user_key, item_key, user_idx, item_idx):
+        #relational neighbor aggregation, h1
+        h1 = self.RelationalAGG(g, self.feature_dict)
+
+        #metapath-based aggregation, h2
+        h2 = {}
+        for key in self.meta_path_patterns.keys():
+            h2[key] = self.hans[key](g, self.feature_dict[key])
+
+        #update node embeddings
+        user_emb = torch.cat((h1[user_key], h2[user_key]), 1)
+        item_emb = torch.cat((h1[item_key], h2[item_key]), 1)
+        user_emb = self.user_layer1(user_emb)
+        item_emb = self.item_layer1(item_emb)
+        user_emb = self.user_layer2(torch.cat((user_emb, self.feature_dict[user_key]), 1))
+        item_emb = self.item_layer2(torch.cat((item_emb, self.feature_dict[item_key]), 1))
+
+        #Relu
+        user_emb = F.relu_(user_emb)
+        item_emb = F.relu_(item_emb)
+        
+        #layer norm
+        user_emb = self.layernorm(user_emb)
+        item_emb = self.layernorm(item_emb)
+        
+        #obtain users/items embeddings and their interactions
+        user_feat = user_emb[user_idx]
+        item_feat = item_emb[item_idx]
+        interaction = user_feat*item_feat
+
+        #score the node pairs
+        pred = self.pred(interaction)
+        pred = self.dropout(pred) #dropout
+        pred = self.fc(pred)
+        pred = torch.sigmoid(pred)
+
+        return pred.squeeze(1)
+        
\ No newline at end of file

examples/pytorch/TAHIN/data_loader.py

+import torch
+from torch.utils.data import Dataset, DataLoader
+import dgl
+import os
+import pickle as pkl
+import numpy as np
+import random
+
+# Split data into train/eval/test
+def split_data(hg, etype_name):
+    src, dst = hg.edges(etype=etype_name)
+    user_item_src = src.numpy().tolist()
+    user_item_dst = dst.numpy().tolist()
+    
+    num_link = len(user_item_src)
+    pos_label=[1]*num_link
+    pos_data=list(zip(user_item_src,user_item_dst,pos_label))
+
+    ui_adj = np.array(hg.adj(etype=etype_name).to_dense())
+    full_idx = np.where(ui_adj==0)
+
+    sample = random.sample(range(0, len(full_idx[0])), num_link)
+    neg_label = [0]*num_link
+    neg_data = list(zip(full_idx[0][sample],full_idx[1][sample],neg_label))
+    
+    full_data = pos_data + neg_data
+    random.shuffle(full_data)
+
+    train_size = int(len(full_data) * 0.6)
+    eval_size = int(len(full_data) * 0.2)
+    test_size = len(full_data) - train_size - eval_size
+    train_data = full_data[:train_size]
+    eval_data = full_data[train_size : train_size+eval_size]
+    test_data = full_data[train_size+eval_size : train_size+eval_size+test_size]
+    train_data = np.array(train_data)
+    eval_data = np.array(eval_data)
+    test_data = np.array(test_data)
+    
+    return train_data, eval_data, test_data
+    
+
+def process_amazon(root_path):
+    # User-Item 3584 2753 50903 UIUI
+    # Item-View 2753 3857 5694 UIVI
+    # Item-Brand 2753 334 2753 UIBI
+    # Item-Category 2753 22 5508 UICI
+    
+    #Construct graph from raw data.
+    # load data of amazon
+    data_path = os.path.join(root_path, 'Amazon')
+    if not (os.path.exists(data_path)):
+        print('Can not find amazon in {}, please download the dataset first.'.format(data_path))
+    
+    # item_view
+    item_view_src=[]
+    item_view_dst=[]
+    with open(os.path.join(data_path, 'item_view.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split(',')
+            item, view= int(_line[0]), int(_line[1])
+            item_view_src.append(item)
+            item_view_dst.append(view)
+
+    # user_item
+    user_item_src=[]
+    user_item_dst=[]
+    with open(os.path.join(data_path, 'user_item.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split('\t')
+            user, item, rate = int(_line[0]), int(_line[1]), int(_line[2])
+            if rate > 3:
+                user_item_src.append(user)
+                user_item_dst.append(item)
+
+    # item_brand
+    item_brand_src=[]
+    item_brand_dst=[]
+    with open(os.path.join(data_path, 'item_brand.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split(',')
+            item, brand= int(_line[0]), int(_line[1])
+            item_brand_src.append(item)
+            item_brand_dst.append(brand)
+
+    # item_category
+    item_category_src=[]
+    item_category_dst=[]
+    with open(os.path.join(data_path, 'item_category.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split(',')
+            item, category= int(_line[0]), int(_line[1])
+            item_category_src.append(item)
+            item_category_dst.append(category)
+
+    #build graph
+    hg = dgl.heterograph({
+        ('item', 'iv', 'view') : (item_view_src, item_view_dst),
+        ('view', 'vi', 'item') : (item_view_dst, item_view_src),
+        ('user', 'ui', 'item') : (user_item_src, user_item_dst),
+        ('item', 'iu', 'user') : (user_item_dst, user_item_src),
+        ('item', 'ib', 'brand') : (item_brand_src, item_brand_dst), 
+        ('brand', 'bi', 'item') : (item_brand_dst, item_brand_src),
+        ('item', 'ic', 'category') : (item_category_src, item_category_dst),
+        ('category', 'ci', 'item') : (item_category_dst, item_category_src)})
+
+    print("Graph constructed.")
+
+    # Split data into train/eval/test
+    train_data, eval_data, test_data = split_data(hg, 'ui')
+
+    #delete the positive edges in eval/test data in the original graph
+    train_pos = np.nonzero(train_data[:,2])
+    train_pos_idx = train_pos[0]
+    user_item_src_processed = train_data[train_pos_idx, 0]
+    user_item_dst_processed = train_data[train_pos_idx, 1]
+    edges_dict = {
+        ('item', 'iv', 'view') : (item_view_src, item_view_dst),
+        ('view', 'vi', 'item') : (item_view_dst, item_view_src),
+        ('user', 'ui', 'item') : (user_item_src_processed, user_item_dst_processed),
+        ('item', 'iu', 'user') : (user_item_dst_processed, user_item_src_processed),
+        ('item', 'ib', 'brand') : (item_brand_src, item_brand_dst), 
+        ('brand', 'bi', 'item') : (item_brand_dst, item_brand_src),
+        ('item', 'ic', 'category') : (item_category_src, item_category_dst),
+        ('category', 'ci', 'item') : (item_category_dst, item_category_src)
+    }
+    nodes_dict = {
+        'user': hg.num_nodes('user'),
+        'item': hg.num_nodes('item'),
+        'view': hg.num_nodes('view'),
+        'brand': hg.num_nodes('brand'),
+        'category': hg.num_nodes('category'),
+    }
+    hg_processed = dgl.heterograph(data_dict = edges_dict, num_nodes_dict = nodes_dict)
+    print("Graph processed.")
+
+    #save the processed data
+    with open(os.path.join(root_path, 'amazon_hg.pkl'), 'wb') as file: 
+        pkl.dump(hg_processed, file)
+    with open(os.path.join(root_path, 'amazon_train.pkl'), 'wb') as file: 
+        pkl.dump(train_data, file)
+    with open(os.path.join(root_path, 'amazon_test.pkl'), 'wb') as file: 
+        pkl.dump(test_data, file)
+    with open(os.path.join(root_path, 'amazon_eval.pkl'), 'wb') as file: 
+        pkl.dump(eval_data, file)
+
+    return hg_processed, train_data, eval_data, test_data
+
+
+def process_movielens(root_path):
+    # User-Movie 943 1682 100000 UMUM
+    # User-Age 943 8 943 UAUM
+    # User-Occupation 943 21 943 UOUM
+    # Movie-Genre 1682 18 2861 UMGM
+
+    data_path = os.path.join(root_path, 'Movielens')
+    if not (os.path.exists(data_path)):
+        print('Can not find movielens in {}, please download the dataset first.'.format(data_path))
+
+    #Construct graph from raw data.
+    # movie_genre
+    movie_genre_src=[]
+    movie_genre_dst=[]
+    with open(os.path.join(data_path, 'movie_genre.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split('\t')
+            movie, genre = int(_line[0]), int(_line[1])
+            movie_genre_src.append(movie)
+            movie_genre_dst.append(genre)
+
+    # user_movie
+    user_movie_src=[]
+    user_movie_dst=[]
+    with open(os.path.join(data_path, 'user_movie.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split('\t')
+            user, item, rate = int(_line[0]), int(_line[1]), int(_line[2])
+            if rate > 3:
+                user_movie_src.append(user)
+                user_movie_dst.append(item)
+
+    # user_occupation
+    user_occupation_src=[]
+    user_occupation_dst=[]
+    with open(os.path.join(data_path, 'user_occupation.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split('\t')
+            user, occupation = int(_line[0]), int(_line[1])
+            user_occupation_src.append(user)
+            user_occupation_dst.append(occupation)
+
+    # user_age
+    user_age_src=[]
+    user_age_dst=[]
+    with open(os.path.join(data_path, 'user_age.dat')) as fin:
+        for line in fin.readlines():
+            _line = line.strip().split('\t')
+            user, age = int(_line[0]), int(_line[1])
+            user_age_src.append(user)
+            user_age_dst.append(age)
+
+    #build graph
+    hg = dgl.heterograph({
+        ('movie', 'mg', 'genre') : (movie_genre_src, movie_genre_dst),
+        ('genre', 'gm', 'movie') : (movie_genre_dst, movie_genre_src),
+        ('user', 'um', 'movie') : (user_movie_src, user_movie_dst),
+        ('movie', 'mu', 'user') : (user_movie_dst, user_movie_src),
+        ('user', 'uo', 'occupation') : (user_occupation_src, user_occupation_dst), 
+        ('occupation', 'ou', 'user') : (user_occupation_dst, user_occupation_src),
+        ('user', 'ua', 'age') : (user_age_src, user_age_dst),
+        ('age', 'au', 'user') : (user_age_dst, user_age_src)})
+
+    print("Graph constructed.")
+
+    # Split data into train/eval/test
+    train_data, eval_data, test_data = split_data(hg, 'um')
+
+    #delete the positive edges in eval/test data in the original graph
+    train_pos = np.nonzero(train_data[:,2])
+    train_pos_idx = train_pos[0]
+    user_movie_src_processed = train_data[train_pos_idx, 0]
+    user_movie_dst_processed = train_data[train_pos_idx, 1]
+    edges_dict = {
+        ('movie', 'mg', 'genre') : (movie_genre_src, movie_genre_dst),
+        ('genre', 'gm', 'movie') : (movie_genre_dst, movie_genre_src),
+        ('user', 'um', 'movie') : (user_movie_src_processed, user_movie_dst_processed),
+        ('movie', 'mu', 'user') : (user_movie_dst_processed, user_movie_src_processed),
+        ('user', 'uo', 'occupation') : (user_occupation_src, user_occupation_dst), 
+        ('occupation', 'ou', 'user') : (user_occupation_dst, user_occupation_src),
+        ('user', 'ua', 'age') : (user_age_src, user_age_dst),
+        ('age', 'au', 'user') : (user_age_dst, user_age_src)
+    }
+    nodes_dict = {
+        'user': hg.num_nodes('user'),
+        'movie': hg.num_nodes('movie'),
+        'genre': hg.num_nodes('genre'),
+        'occupation': hg.num_nodes('occupation'),
+        'age': hg.num_nodes('age'),
+    }
+    hg_processed = dgl.heterograph(data_dict = edges_dict, num_nodes_dict = nodes_dict)
+    print("Graph processed.")
+
+    #save the processed data
+    with open(os.path.join(root_path, 'movielens_hg.pkl'), 'wb') as file: 
+        pkl.dump(hg_processed, file)
+    with open(os.path.join(root_path, 'movielens_train.pkl'), 'wb') as file: 
+        pkl.dump(train_data, file)
+    with open(os.path.join(root_path, 'movielens_test.pkl'), 'wb') as file: 
+        pkl.dump(test_data, file)
+    with open(os.path.join(root_path, 'movielens_eval.pkl'), 'wb') as file: 
+        pkl.dump(eval_data, file)
+
+    return hg_processed, train_data, eval_data, test_data
+
+class MyDataset(Dataset):
+
+    def __init__(self, triple):
+
+        self.triple = triple
+        self.len = self.triple.shape[0]
+    
+    def __getitem__(self, index):
+        return self.triple[index, 0], self.triple[index, 1], self.triple[index, 2].float()
+
+    def __len__(self):
+        return self.len
+
+def load_data(dataset, batch_size=128, num_workers = 10, root_path = './data'):
+    if (os.path.exists(os.path.join(root_path, dataset+'_train.pkl'))):
+        g_file = open(os.path.join(root_path, dataset+'_hg.pkl'), 'rb')
+        hg = pkl.load(g_file)
+        g_file.close()
+        train_set_file = open(os.path.join(root_path, dataset+'_train.pkl'), 'rb')
+        train_set = pkl.load(train_set_file)
+        train_set_file.close()
+        test_set_file = open(os.path.join(root_path, dataset+'_test.pkl'), 'rb')
+        test_set = pkl.load(test_set_file)
+        test_set_file.close()
+        eval_set_file = open(os.path.join(root_path, dataset+'_eval.pkl'), 'rb')
+        eval_set = pkl.load(eval_set_file)
+        eval_set_file.close()
+    else: 
+        if dataset == 'movielens':
+            hg, train_set, eval_set, test_set = process_movielens(root_path)
+        elif dataset == 'amazon':
+            hg, train_set, eval_set, test_set = process_amazon(root_path)
+        else:
+            print('Available datasets: movielens, amazon.')
+            raise NotImplementedError
+
+    if dataset == 'movielens':
+        meta_paths = {
+            'user': [['um', 'mu']], 
+            'movie': [['mu', 'um'], ['mg', 'gm']]
+        }
+        user_key = 'user'
+        item_key = 'movie'
+    elif dataset == 'amazon':
+        meta_paths = {
+            'user': [['ui', 'iu']],
+            'item': [['iu', 'ui'], ['ic', 'ci'], ['ib', 'bi'], ['iv', 'vi']]
+        }
+        user_key = 'user'
+        item_key = 'item'
+    else:
+        print('Available datasets: movielens, amazon.')
+        raise NotImplementedError
+    
+    train_set = torch.Tensor(train_set).long()
+    eval_set = torch.Tensor(eval_set).long()
+    test_set = torch.Tensor(test_set).long()
+    
+    train_set = MyDataset(train_set)
+    train_loader= DataLoader(dataset=train_set, batch_size = batch_size, shuffle=True, num_workers = num_workers)
+    eval_set = MyDataset(eval_set)
+    eval_loader= DataLoader(dataset=eval_set, batch_size = batch_size, shuffle=True, num_workers = num_workers)
+    test_set = MyDataset(test_set)
+    test_loader= DataLoader(dataset=test_set, batch_size = batch_size, shuffle=True, num_workers = num_workers)
+
+    return hg, train_loader, eval_loader, test_loader, meta_paths, user_key, item_key
+    

examples/pytorch/TAHIN/main.py

+import torch 
+import torch.nn as nn
+import torch.optim as optim
+import dgl
+import numpy as np
+import pickle as pkl
+import argparse
+from data_loader import load_data
+from TAHIN import TAHIN
+from utils import evaluate_auc, evaluate_acc, evaluate_f1_score, evaluate_logloss
+
+def main(args):
+    #step 1: Check device
+    if args.gpu >= 0 and torch.cuda.is_available():
+        device = 'cuda:{}'.format(args.gpu)
+    else:
+        device = 'cpu'
+
+    #step 2: Load data 
+    g, train_loader, eval_loader, test_loader, meta_paths, user_key, item_key = load_data(args.dataset, args.batch, args.num_workers, args.path)
+    g = g.to(device)
+    print('Data loaded.')
+
+    #step 3: Create model and training components
+    model = TAHIN(
+        g, meta_paths, args.in_size, args.out_size, args.num_heads, args.dropout
+        )
+    model = model.to(device)
+    criterion = nn.BCELoss()
+    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)
+    print('Model created.')
+
+    #step 4: Training
+    print('Start training.')
+    best_acc = 0.0
+    kill_cnt = 0
+    for epoch in range(args.epochs):
+        # Training and validation using a full graph
+        model.train()
+        train_loss = []
+        for step, batch in enumerate(train_loader):
+            user, item, label = [_.to(device) for _ in batch]
+            logits = model.forward(g, user_key, item_key, user, item)
+
+            # compute loss
+            tr_loss = criterion(logits, label)
+            train_loss.append(tr_loss)
+
+            # backward
+            optimizer.zero_grad()
+            tr_loss.backward()
+            optimizer.step()
+
+        train_loss = np.sum(train_loss)
+
+        model.eval()
+        with torch.no_grad():
+            validate_loss = []
+            validate_acc = []
+            for step, batch in enumerate(eval_loader):
+                user, item, label = [_.to(device) for _ in batch]
+                logits = model.forward(g, user_key, item_key, user, item)
+
+                # compute loss
+                val_loss = criterion(logits, label)
+                val_acc = evaluate_acc(logits.detach().cpu().numpy(), label.detach().cpu().numpy())
+                validate_loss.append(val_loss)
+                validate_acc.append(val_acc)
+        
+            validate_loss = np.sum(validate_loss)
+            validate_acc = np.mean(validate_acc)
+        
+            #validate
+            if validate_acc > best_acc:
+                best_acc = validate_acc
+                best_epoch = epoch
+                torch.save(model.state_dict(), 'TAHIN'+'_'+args.dataset)
+                kill_cnt = 0
+                print("saving model...")
+            else:
+                kill_cnt += 1
+                if kill_cnt > args.early_stop:
+                    print('early stop.')
+                    print("best epoch:{}".format(best_epoch))
+                    break
+
+            print("In epoch {}, Train Loss: {:.4f}, Valid Loss: {:.5}\n, Valid ACC: {:.5}".format(epoch, train_loss, validate_loss, validate_acc))
+
+    #test use the best model
+    model.eval()
+    with torch.no_grad():
+        model.load_state_dict(torch.load('TAHIN'+'_'+args.dataset))
+        test_loss = []
+        test_acc = []
+        test_auc = []
+        test_f1 = []
+        test_logloss = []
+        for step, batch in enumerate(test_loader):
+            user, item, label = [_.to(device) for _ in batch]
+            logits = model.forward(g, user_key, item_key, user, item)
+
+            # compute loss
+            loss = criterion(logits, label)
+            acc = evaluate_acc(logits.detach().cpu().numpy(), label.detach().cpu().numpy())
+            auc = evaluate_auc(logits.detach().cpu().numpy(), label.detach().cpu().numpy())
+            f1 = evaluate_f1_score(logits.detach().cpu().numpy(), label.detach().cpu().numpy())
+            log_loss = evaluate_logloss(logits.detach().cpu().numpy(), label.detach().cpu().numpy())
+            
+            test_loss.append(loss)
+            test_acc.append(acc)
+            test_auc.append(auc)
+            test_f1.append(f1)
+            test_logloss.append(log_loss)
+        
+        test_loss = np.sum(test_loss)
+        test_acc = np.mean(test_acc)
+        test_auc = np.mean(test_auc)
+        test_f1 = np.mean(test_f1)
+        test_logloss = np.mean(test_logloss)
+        print("Test Loss: {:.5}\n, Test ACC: {:.5}\n, AUC: {:.5}\n, F1: {:.5}\n, Logloss: {:.5}\n".format(test_loss, test_acc, test_auc, test_f1, test_logloss))
+    
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Parser For Arguments', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    parser.add_argument('--dataset', default='movielens', help='Dataset to use, default: movielens')
+    parser.add_argument('--path', default='./data', help='Path to save the data')
+    parser.add_argument('--model', default='TAHIN', help='Model Name')
+
+    parser.add_argument('--batch', default=128, type=int, help='Batch size')
+    parser.add_argument('--gpu', type=int, default='0', help='Set GPU Ids : Eg: For CPU = -1, For Single GPU = 0')
+    parser.add_argument('--epochs', type=int, default=500, help='Maximum number of epochs')
+    parser.add_argument('--wd', type=float, default=0, help='L2 Regularization for Optimizer')
+    parser.add_argument('--lr', type=float, default=0.001, help='Learning Rate')
+    parser.add_argument('--num_workers', type=int, default=10, help='Number of processes to construct batches')
+    parser.add_argument('--early_stop', default=15, type=int, help='Patience for early stop.')
+    
+    parser.add_argument('--in_size', default=128, type=int, help='Initial dimension size for entities.')
+    parser.add_argument('--out_size', default=128, type=int, help='Output dimension size for entities.')
+    
+    parser.add_argument('--num_heads', default=1, type=int, help='Number of attention heads')
+    parser.add_argument('--dropout', default=0.1, type=float, help='Dropout.')
+    
+    args = parser.parse_args()
+
+    print(args)
+
+    main(args)
+    
\ No newline at end of file

examples/pytorch/TAHIN/readme.md

+# DGL Implementation of the TAHIN
+
+This DGL example implements the TAHIN module proposed in the paper [HCDIR](https://arxiv.org/pdf/2007.15293.pdf). Since the code and dataset have not been published yet, we implement its main idea and experiment on two other datasets.
+
+Example implementor
+----------------------
+This example was implemented by [KounianhuaDu](https://github.com/KounianhuaDu) during her software development intern time at the AWS Shanghai AI Lab.
+
+Dependencies
+----------------------
+- pytorch 1.7.1
+- dgl 0.6.0
+
+Datasets
+---------------------------------------
+The datasets used can be downloaded from [here](https://github.com/librahu/HIN-Datasets-for-Recommendation-and-Network-Embedding). For the experiments, all the positive edges are fetched and the same number of negative edges are randomly sampled. The edges are then shuffled and splitted into train/validate/test at a ratio of 6:2:2. The positive edges that appear in the validation and test sets are then removed from the original graph.
+
+The original graph statistics:
+
+**Movielens** 
+
+(Source : https://grouplens.org/datasets/movielens/)
+
+| Entity         |#Entity        |
+| :-------------:|:-------------:|
+| User           | 943           |
+| Age            | 8             |
+| Occupation     | 21            |
+| Movie          | 1,682         |
+| Genre          | 18            |
+
+| Relation            |#Relation      |
+| :-------------:     |:-------------:|
+| User - Movie        | 100,000       |
+| User - User (KNN)   | 47,150        |
+| User - Age          | 943           |
+| User - Occupation   | 943           |
+| Movie - Movie (KNN) | 82,798        |
+| Movie - Genre       | 2,861         |
+
+**Amazon** 
+
+(Source : http://jmcauley.ucsd.edu/data/amazon/)
+
+| Entity         |#Entity        |
+| :-------------:|:-------------:|
+| User           | 6,170         |
+| Item           | 2,753         |
+| View           | 3,857         |           
+| Category       | 22            |
+| Brand          | 334           |
+
+| Relation          |#Relation      |
+| :-------------:   |:-------------:|
+| User - Item       | 195,791       |
+| Item - View       | 5,694         |
+| Item - Category   | 5,508         | 
+| Item - Brand      | 2,753         |
+
+How to run
+--------------------------------
+
+```python
+python main.py --dataset amazon --gpu 0
+```
+
+
+```python
+python main.py --dataset movielens --gpu 0
+```
+
+
+Performance
+-------------------------
+**Results**
+
+| Dataset |         Movielens        |          Amazon          |
+|---------| ------------------------ | ------------------------ |
+|  Metric |    HAN     /    TAHIN    |    HAN     /    TAHIN    |
+|   AUC   |   0.9297   /   0.9392    |   0.8470   /   0.8442    |
+|   ACC   |   0.8627   /   0.8683    |   0.7672   /   0.7619    |
+|    F1   |   0.8631   /   0.8707    |   0.7628   /   0.7499    |
+| Logloss |   0.3689   /   0.3266    |   0.5311   /   0.5150    |

examples/pytorch/TAHIN/utils.py

+import numpy as np
+import torch
+from sklearn.metrics import roc_auc_score, accuracy_score, log_loss, f1_score, average_precision_score, ndcg_score
+
+def evaluate_auc(pred, label):
+    res=roc_auc_score(y_score=pred, y_true=label)
+    return res
+
+def evaluate_acc(pred, label):
+    res = []
+    for _value in pred:
+        if _value >= 0.5:
+            res.append(1)
+        else:
+            res.append(0)
+    return accuracy_score(y_pred=res, y_true=label)
+
+def evaluate_f1_score(pred, label):
+    res = []
+    for _value in pred:
+        if _value >= 0.5:
+            res.append(1)
+        else:
+            res.append(0)
+    return f1_score(y_pred=res, y_true=label)
+
+def evaluate_logloss(pred, label):
+    res = log_loss(y_true=label, y_pred=pred,eps=1e-7, normalize=True)
+    return res
+    
\ No newline at end of file