数据处理完整代码

到这里，我们已完成了ml-1m数据读取和处理，接下来，我们将数据处理的代码封装到一个Python类中，完整实现如下：

  1. import random
  2. import numpy as np
  3. from PIL import Image
  5. class MovieLen(object):
  6.     def __init__(self, use_poster):
  7.         self.use_poster = use_poster
  8.         # 声明每个数据文件的路径
  9.         usr_info_path = "./work/ml-1m/users.dat"
  10.         if use_poster:
  11.             rating_path = "./work/ml-1m/new_rating.txt"
  12.         else:
  13.             rating_path = "./work/ml-1m/ratings.dat"
  15.         movie_info_path = "./work/ml-1m/movies.dat"
  16.         self.poster_path = "./work/ml-1m/posters/"
  17.         # 得到电影数据
  18.         self.movie_info, self.movie_cat, self.movie_title = self.get_movie_info(movie_info_path)
  19.         # 记录电影的最大ID
  20.         self.max_mov_cat = np.max([self.movie_cat[k] for k in self.movie_cat])
  21.         self.max_mov_tit = np.max([self.movie_title[k] for k in self.movie_title])
  22.         self.max_mov_id = np.max(list(map(int, self.movie_info.keys())))
  23.         # 记录用户数据的最大ID
  24.         self.max_usr_id = 0
  25.         self.max_usr_age = 0
  26.         self.max_usr_job = 0
  27.         # 得到用户数据
  28.         self.usr_info = self.get_usr_info(usr_info_path)
  29.         # 得到评分数据
  30.         self.rating_info = self.get_rating_info(rating_path)
  31.         # 构建数据集 
  32.         self.dataset = self.get_dataset(usr_info=self.usr_info,
  33.                                         rating_info=self.rating_info,
  34.                                         movie_info=self.movie_info)
  35.         # 划分数据集，获得数据加载器
  36.         self.train_dataset = self.dataset[:int(len(self.dataset)*0.9)]
  37.         self.valid_dataset = self.dataset[int(len(self.dataset)*0.9):]
  38.         print("##Total dataset instances: ", len(self.dataset))
  39.         print("##MovieLens dataset information: \nusr num: {}\n"
  40.               "movies num: {}".format(len(self.usr_info),len(self.movie_info)))
  41.     # 得到电影数据
  42.     def get_movie_info(self, path):
  43.         # 打开文件，编码方式选择ISO-8859-1，读取所有数据到data中 
  44.         with open(path, 'r', encoding="ISO-8859-1") as f:
  45.             data = f.readlines()
  46.         # 建立三个字典，分别用户存放电影所有信息，电影的名字信息、类别信息
  47.         movie_info, movie_titles, movie_cat = {}, {}, {}
  48.         # 对电影名字、类别中不同的单词计数
  49.         t_count, c_count = 1, 1
  51.         count_tit = {}
  52.         # 按行读取数据并处理
  53.         for item in data:
  54.             item = item.strip().split("::")
  55.             v_id = item[0]
  56.             v_title = item[1][:-7]
  57.             cats = item[2].split('|')
  58.             v_year = item[1][-5:-1]
  60.             titles = v_title.split()
  61.             # 统计电影名字的单词，并给每个单词一个序号，放在movie_titles中
  62.             for t in titles:
  63.                 if t not in movie_titles:
  64.                     movie_titles[t] = t_count
  65.                     t_count += 1
  66.             # 统计电影类别单词，并给每个单词一个序号，放在movie_cat中
  67.             for cat in cats:
  68.                 if cat not in movie_cat:
  69.                     movie_cat[cat] = c_count
  70.                     c_count += 1
  71.             # 补0使电影名称对应的列表长度为15
  72.             v_tit = [movie_titles[k] for k in titles]
  73.             while len(v_tit)<15:
  74.                 v_tit.append(0)
  75.             # 补0使电影种类对应的列表长度为6
  76.             v_cat = [movie_cat[k] for k in cats]
  77.             while len(v_cat)<6:
  78.                 v_cat.append(0)
  79.             # 保存电影数据到movie_info中
  80.             movie_info[v_id] = {'mov_id': int(v_id),
  81.                                 'title': v_tit,
  82.                                 'category': v_cat,
  83.                                 'years': int(v_year)}
  84.         return movie_info, movie_cat, movie_titles
  86.     def get_usr_info(self, path):
  87.         # 性别转换函数，M-0， F-1
  88.         def gender2num(gender):
  89.             return 1 if gender == 'F' else 0
  91.         # 打开文件，读取所有行到data中
  92.         with open(path, 'r') as f:
  93.             data = f.readlines()
  94.         # 建立用户信息的字典
  95.         use_info = {}
  97.         max_usr_id = 0
  98.         #按行索引数据
  99.         for item in data:
  100.             # 去除每一行中和数据无关的部分
  101.             item = item.strip().split("::")
  102.             usr_id = item[0]
  103.             # 将字符数据转成数字并保存在字典中
  104.             use_info[usr_id] = {'usr_id': int(usr_id),
  105.                                 'gender': gender2num(item[1]),
  106.                                 'age': int(item[2]),
  107.                                 'job': int(item[3])}
  108.             self.max_usr_id = max(self.max_usr_id, int(usr_id))
  109.             self.max_usr_age = max(self.max_usr_age, int(item[2]))
  110.             self.max_usr_job = max(self.max_usr_job, int(item[3]))
  111.         return use_info
  112.     # 得到评分数据
  113.     def get_rating_info(self, path):
  114.         # 读取文件里的数据
  115.         with open(path, 'r') as f:
  116.             data = f.readlines()
  117.         # 将数据保存在字典中并返回
  118.         rating_info = {}
  119.         for item in data:
  120.             item = item.strip().split("::")
  121.             usr_id,movie_id,score = item[0],item[1],item[2]
  122.             if usr_id not in rating_info.keys():
  123.                 rating_info[usr_id] = {movie_id:float(score)}
  124.             else:
  125.                 rating_info[usr_id][movie_id] = float(score)
  126.         return rating_info
  127.     # 构建数据集
  128.     def get_dataset(self, usr_info, rating_info, movie_info):
  129.         trainset = []
  130.         for usr_id in rating_info.keys():
  131.             usr_ratings = rating_info[usr_id]
  132.             for movie_id in usr_ratings:
  133.                 trainset.append({'usr_info': usr_info[usr_id],
  134.                                  'mov_info': movie_info[movie_id],
  135.                                  'scores': usr_ratings[movie_id]})
  136.         return trainset
  138.     def load_data(self, dataset=None, mode='train'):
  139.         use_poster = False
  141.         # 定义数据迭代Batch大小
  142.         BATCHSIZE = 256
  144.         data_length = len(dataset)
  145.         index_list = list(range(data_length))
  146.         # 定义数据迭代加载器
  147.         def data_generator():
  148.             # 训练模式下，打乱训练数据
  149.             if mode == 'train':
  150.                 random.shuffle(index_list)
  151.             # 声明每个特征的列表
  152.             usr_id_list,usr_gender_list,usr_age_list,usr_job_list = [], [], [], []
  153.             mov_id_list,mov_tit_list,mov_cat_list,mov_poster_list = [], [], [], []
  154.             score_list = []
  155.             # 索引遍历输入数据集
  156.             for idx, i in enumerate(index_list):
  157.                 # 获得特征数据保存到对应特征列表中
  158.                 usr_id_list.append(dataset[i]['usr_info']['usr_id'])
  159.                 usr_gender_list.append(dataset[i]['usr_info']['gender'])
  160.                 usr_age_list.append(dataset[i]['usr_info']['age'])
  161.                 usr_job_list.append(dataset[i]['usr_info']['job'])
  163.                 mov_id_list.append(dataset[i]['mov_info']['mov_id'])
  164.                 mov_tit_list.append(dataset[i]['mov_info']['title'])
  165.                 mov_cat_list.append(dataset[i]['mov_info']['category'])
  166.                 mov_id = dataset[i]['mov_info']['mov_id']
  168.                 if use_poster:
  169.                     # 不使用图像特征时，不读取图像数据，加快数据读取速度
  170.                     poster = Image.open(self.poster_path+'mov_id{}.jpg'.format(str(mov_id[0])))
  171.                     poster = poster.resize([64, 64])
  172.                     if len(poster.size) <= 2:
  173.                         poster = poster.convert("RGB")
  175.                     mov_poster_list.append(np.array(poster))
  177.                 score_list.append(int(dataset[i]['scores']))
  178.                 # 如果读取的数据量达到当前的batch大小，就返回当前批次
  179.                 if len(usr_id_list)==BATCHSIZE:
  180.                     # 转换列表数据为数组形式，reshape到固定形状
  181.                     usr_id_arr = np.array(usr_id_list)
  182.                     usr_gender_arr = np.array(usr_gender_list)
  183.                     usr_age_arr = np.array(usr_age_list)
  184.                     usr_job_arr = np.array(usr_job_list)
  186.                     mov_id_arr = np.array(mov_id_list)
  187.                     mov_cat_arr = np.reshape(np.array(mov_cat_list), [BATCHSIZE, 6]).astype(np.int64)
  188.                     mov_tit_arr = np.reshape(np.array(mov_tit_list), [BATCHSIZE, 1, 15]).astype(np.int64)
  190.                     if use_poster:
  191.                         mov_poster_arr = np.reshape(np.array(mov_poster_list)/127.5 - 1, [BATCHSIZE, 3, 64, 64]).astype(np.float32)
  192.                     else:
  193.                         mov_poster_arr = np.array([0.])
  195.                     scores_arr = np.reshape(np.array(score_list), [-1, 1]).astype(np.float32)
  197.                     # 放回当前批次数据
  198.                     yield [usr_id_arr, usr_gender_arr, usr_age_arr, usr_job_arr], \
  199.                            [mov_id_arr, mov_cat_arr, mov_tit_arr, mov_poster_arr], scores_arr
  201.                     # 清空数据
  202.                     usr_id_list, usr_gender_list, usr_age_list, usr_job_list = [], [], [], []
  203.                     mov_id_list, mov_tit_list, mov_cat_list, score_list = [], [], [], []
  204.                     mov_poster_list = []
  205.         return data_generator
  207. # 声明数据读取类
  208. dataset = MovieLen(False)
  209. # 定义数据读取器
  210. train_loader = dataset.load_data(dataset=dataset.train_dataset, mode='train')
  211. # 迭代的读取数据， Batchsize = 256
  212. for idx, data in enumerate(train_loader()):
  213.     usr, mov, score = data
  214.     print("打印用户ID，性别，年龄，职业数据的维度：")
  215.     for v in usr:
  216.         print(v.shape)
  217.     print("打印电影ID，名字，类别数据的维度：")
  218.     for v in mov:
  219.         print(v.shape)
  221.     break
  1. ##Total dataset instances:  1000209
  2. ##MovieLens dataset information: 
  3. usr num: 6040
  4. movies num: 3883
  5. 打印用户ID，性别，年龄，职业数据的维度：
  6. (256,)
  7. (256,)
  8. (256,)
  9. (256,)
  10. 打印电影ID，名字，类别数据的维度：
  11. (256,)
  12. (256, 6)
  13. (256, 1, 15)
  14. (1,)