利用电影评分数据集
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
data = pd.read_csv('movie/ratings.csv')
data_train = data.pivot(index = 'userId', columns = 'movieId', values = 'rating')
matrax = data_train.fillna(0)
matrix = np.array(matrax)
data = torch.tensor(matrix)
u, s, v = torch.svd(data)
s_topk = torch.topk(s, 100)
idx_list = s_topk[1].tolist()
u_topk = torch.index_select(u, 1, s_topk[1].squeeze())
v_topk = torch.index_select(v, 1, s_topk[1].squeeze())
matrix_data = torch.mm(u_topk, v_topk.t())