Pytorch-SVD分解

利用电影评分数据集

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)
#取前100个特征
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())