使用Movielens-20m创建DGL异质图

数据：

code

def add_MG():
    """ 根据MovieID_genreID.csv创建MG边 """
    # 当加载csv文件的多列数据时可以使用unpack将加载的数据列进场解耦到不同数组中
    filename = "data/Movielens-20m/MovieID_genreID.csv"
    movies, genres = np.loadtxt(filename, delimiter=",", skiprows=1, usecols=(0, 1), unpack=True, dtype=int)
    movies, genres = movies.tolist(), genres.tolist()
    MG_data, GM_data = (movies, genres), (genres, movies)
    return MG_data, GM_data


def add_UMR():
    filename = 'data/Movielens-20m/UserID_MovieID_Rating.csv'
    users, movies, ratings = np.loadtxt(filename, delimiter=",", skiprows=1, usecols=(0, 1, 2), unpack=True)
    users, movies, ratings = users.astype(dtype=int).tolist(), movies.astype(dtype=int).tolist(), ratings.tolist()
    UMR_data, MUR_data = (users, movies), (movies, users)
    return UMR_data, MUR_data, ratings

def add_UMT():
    filename = 'data/Movielens-20m/UserID_MovieID_TagID.csv'
    users, movies, tags = np.loadtxt(filename, delimiter=",", skiprows=1, usecols=(0, 1, 2), unpack=True, dtype=int)
    users, movies, tags = users.tolist(), movies.tolist(), tags.tolist()
    UMT_data, MUT_data = (users, movies), (movies, users)
    return UMT_data, MUT_data, tags

def createHeteroGraph():
    MG_data, GM_data = add_MG()
    UMR_data, MUR_data, ratings = add_UMR()
    UMT_data, MUT_data, tags = add_UMT()
    g = dgl.heterograph({
        ('movie', 'MG', 'genre'): MG_data,
        ('genre', 'GM', 'movie'): GM_data,
        ('user', 'UMR', 'movie'): UMR_data,
        ('movie', 'MUR', 'user'): MUR_data,
        ('user', 'UMT', 'movie'): UMT_data,
        ('movie', 'MUT', 'user'): MUT_data
    })
    g.edges['UMR'].data['rate'] = th.tensor(ratings)
    g.edges['MUR'].data['rate'] = th.tensor(ratings)
    g.edges['UMT'].data['tag'] = th.tensor(tags)
    g.edges['MUT'].data['tag'] = th.tensor(tags)
    return g

if __name__ == '__main__':
    g = createHeteroGraph()
    print(g)

output

D:\Apps\Anaconda3\python.exe D:/PYproject/DataSet_Process/MovieLens-20m.py
Using backend: pytorch
Graph(num_nodes={'genre': 19, 'movie': 131263, 'user': 138494},
      num_edges={('genre', 'GM', 'movie'): 54160, ('movie', 'MG', 'genre'): 54160, ('movie', 'MUR', 'user'): 20000263, ('movie', 'MUT', 'user'): 217571, ('user', 'UMR', 'movie'): 20000263, ('user', 'UMT', 'movie'): 217571},
      metagraph=[('genre', 'movie', 'GM'), ('movie', 'genre', 'MG'), ('movie', 'user', 'MUR'), ('movie', 'user', 'MUT'), ('user', 'movie', 'UMR'), ('user', 'movie', 'UMT')])

Process finished with exit code 0