SSLRec代码分析

encoder-models-general_cf

autocf.py

import torch as t  # 导入PyTorch并重命名为t
from torch import nn  # 从PyTorch导入神经网络模块
import torch.nn.functional as F  # 导入PyTorch中的函数式API
from config.configurator import configs  # 从配置模块导入configs
from models.loss_utils import reg_params  # 从模型损失工具中导入reg_params
from models.base_model import BaseModel  # 从基模型模块导入BaseModel# 初始化一些常用的权重初始化方法
init = nn.init.xavier_uniform_
uniformInit = nn.init.uniform# 定义AutoCF类，继承自BaseModel
class AutoCF(BaseModel):def __init__(self, data_handler):super(AutoCF, self).__init__(data_handler)# 定义用户和物品的嵌入矩阵self.user_embeds = nn.Parameter(init(t.empty(self.user_num, self.embedding_size)))self.item_embeds = nn.Parameter(init(t.empty(self.item_num, self.embedding_size)))self.adj = data_handler.torch_adj  # 获取邻接矩阵self.all_one_adj = self.make_all_one_adj()  # 创建全1的邻接矩阵# 获取超参数配置self.gt_layer = configs['model']['gt_layer']self.gcn_layer = self.hyper_config['gcn_layer']self.reg_weight = self.hyper_config['reg_weight']self.ssl_reg = self.hyper_config['ssl_reg']# 定义GCN层和GT层self.gcnLayers = nn.Sequential(*[GCNLayer() for i in range(self.gcn_layer)])self.gtLayers = nn.Sequential(*[GTLayer() for i in range(self.gt_layer)])self.masker = RandomMaskSubgraphs()  # 随机掩码子图self.sampler = LocalGraph()  # 局部图采样def make_all_one_adj(self):idxs = self.adj._indices()  # 获取稀疏矩阵的索引vals = t.ones_like(self.adj._values())  # 将值全设置为1shape = self.adj.shape  # 获取矩阵的形状return t.sparse.FloatTensor(idxs, vals, shape).cuda()  # 创建稀疏矩阵并移动到CUDAdef get_ego_embeds(self):return t.concat([self.user_embeds, self.item_embeds], axis=0)  # 合并用户和物品的嵌入矩阵def sample_subgraphs(self):return self.sampler(self.all_one_adj, self.get_ego_embeds())  # 采样子图def mask_subgraphs(self, seeds):return self.masker(self.adj, seeds)  # 对子图进行掩码def forward(self, encoder_adj, decoder_adj=None):embeds = t.concat([self.user_embeds, self.item_embeds], axis=0)  # 合并嵌入矩阵embedsLst = [embeds]  # 初始化嵌入矩阵列表for i, gcn in enumerate(self.gcnLayers):embeds = gcn(encoder_adj, embedsLst[-1])  # 通过GCN层传播embedsLst.append(embeds)if decoder_adj is not None:for gt in self.gtLayers:embeds = gt(decoder_adj, embedsLst[-1])  # 通过GT层传播embedsLst.append(embeds)embeds = sum(embedsLst)  # 求和所有层的嵌入矩阵return embeds[:self.user_num], embeds[self.user_num:]  # 返回用户和物品的嵌入矩阵def contrast(self, nodes, allEmbeds, allEmbeds2=None):if allEmbeds2 is not None:pckEmbeds = allEmbeds[nodes]scores = t.log(t.exp(pckEmbeds @ allEmbeds2.T).sum(-1)).mean()else:uniqNodes = t.unique(nodes)pckEmbeds = allEmbeds[uniqNodes]scores = t.log(t.exp(pckEmbeds @ allEmbeds.T).sum(-1)).mean()return scoresdef cal_loss(self, batch_data, encoder_adj, decoder_adj):user_embeds, item_embeds = self.forward(encoder_adj, decoder_adj)  # 获取前向传播后的嵌入矩阵ancs, poss, _ = batch_data  # 解包批处理数据anc_embeds = user_embeds[ancs]  # 获取锚点用户的嵌入pos_embeds = item_embeds[poss]  # 获取正样本物品的嵌入rec_loss = (-t.sum(anc_embeds * pos_embeds, dim=-1)).mean()  # 计算推荐损失reg_loss = reg_params(self) * self.reg_weight  # 计算正则化损失cl_loss = (self.contrast(ancs, user_embeds) + self.contrast(poss, item_embeds)) * self.ssl_reg + self.contrast(ancs, user_embeds, item_embeds)  # 计算对比学习损失loss = rec_loss + reg_loss + cl_loss  # 总损失losses = {'rec_loss': rec_loss, 'reg_loss': reg_loss, 'cl_loss': cl_loss}  # 各种损失的字典return loss, losses  # 返回总损失和各个损失def full_predict(self, batch_data):user_embeds, item_embeds = self.forward(self.adj, self.adj)  # 前向传播pck_users, train_mask = batch_data  # 解包批处理数据pck_users = pck_users.long()  # 转换为长整型pck_user_embeds = user_embeds[pck_users]  # 获取选择用户的嵌入full_preds = pck_user_embeds @ item_embeds.T  # 计算预测分数full_preds = self._mask_predict(full_preds, train_mask)  # 掩码预测return full_preds  # 返回预测分数# 定义GCN层
class GCNLayer(nn.Module):def __init__(self):super(GCNLayer, self).__init__()def forward(self, adj, embeds):return t.spmm(adj, embeds)  # 稀疏矩阵乘法# 定义GT层
class GTLayer(nn.Module):def __init__(self):super(GTLayer, self).__init__()self.head_num = configs['model']['head_num']  # 获取头的数量self.embedding_size = configs['model']['embedding_size']  # 获取嵌入矩阵的大小self.qTrans = nn.Parameter(init(t.empty(self.embedding_size, self.embedding_size)))  # 初始化查询转换矩阵self.kTrans = nn.Parameter(init(t.empty(self.embedding_size, self.embedding_size)))  # 初始化键转换矩阵self.vTrans = nn.Parameter(init(t.empty(self.embedding_size, self.embedding_size)))  # 初始化值转换矩阵def forward(self, adj, embeds):indices = adj._indices()  # 获取稀疏矩阵的索引rows, cols = indices[0, :], indices[1, :]  # 获取行和列的索引rowEmbeds = embeds[rows]  # 获取行嵌入colEmbeds = embeds[cols]  # 获取列嵌入qEmbeds = (rowEmbeds @ self.qTrans).view([-1, self.head_num, self.embedding_size // self.head_num])  # 计算查询嵌入kEmbeds = (colEmbeds @ self.kTrans).view([-1, self.head_num, self.embedding_size // self.head_num])  # 计算键嵌入vEmbeds = (colEmbeds @ self.vTrans).view([-1, self.head_num, self.embedding_size // self.head_num])  # 计算值嵌入att = t.einsum('ehd, ehd -> eh', qEmbeds, kEmbeds)  # 计算注意力权重att = t.clamp(att, -10.0, 10.0)  # 截断权重expAtt = t.exp(att)  # 计算指数权重tem = t.zeros([adj.shape[0], self.head_num]).cuda()attNorm = (tem.index_add_(0, rows, expAtt))[rows]att = expAtt / (attNorm + 1e-8)  # 归一化注意力权重resEmbeds = t.einsum('eh, ehd -> ehd', att, vEmbeds).view([-1, self.embedding_size])  # 计算结果嵌入tem = t.zeros([adj.shape[0], self.embedding_size]).cuda()resEmbeds = tem.index_add_(0, rows, resEmbeds)  # 累加结果嵌入return resEmbeds  # 返回结果
class LocalGraph(nn.Module):def __init__(self):super(LocalGraph, self).__init__()self.seed_num = configs['model']['seed_num']  # 从配置中获取种子节点数量def makeNoise(self, scores):noise = t.rand(scores.shape).cuda()  # 生成与得分相同形状的随机噪声noise[noise == 0] = 1e-8  # 防止噪声为0的情况noise = -t.log(-t.log(noise))  # 双对数变换return t.log(scores) + noise  # 返回加噪声后的得分def forward(self, allOneAdj, embeds):# allOneAdj应为无自环的邻接矩阵# embeds应为零阶嵌入order = t.sparse.sum(allOneAdj, dim=-1).to_dense().view([-1, 1])  # 计算节点的度并转为密集矩阵fstEmbeds = t.spmm(allOneAdj, embeds) - embeds  # 计算第一层嵌入fstNum = order  # 第一层的节点度scdEmbeds = (t.spmm(allOneAdj, fstEmbeds) - fstEmbeds) - order * embeds  # 计算第二层嵌入scdNum = (t.spmm(allOneAdj, fstNum) - fstNum) - order  # 计算第二层的节点度subgraphEmbeds = (fstEmbeds + scdEmbeds) / (fstNum + scdNum + 1e-8)  # 计算子图嵌入subgraphEmbeds = F.normalize(subgraphEmbeds, p=2)  # 对子图嵌入进行归一化embeds = F.normalize(embeds, p=2)  # 对嵌入进行归一化scores = t.sigmoid(t.sum(subgraphEmbeds * embeds, dim=-1))  # 计算得分scores = self.makeNoise(scores)  # 添加噪声_, seeds = t.topk(scores, self.seed_num)  # 获取得分最高的种子节点return scores, seeds  # 返回得分和种子节点class RandomMaskSubgraphs(nn.Module):def __init__(self):super(RandomMaskSubgraphs, self).__init__()self.flag = False  # 初始化标志位self.mask_depth = configs['model']['mask_depth']  # 获取掩码深度self.keep_rate = configs['model']['keep_rate']  # 获取保留率self.user_num = configs['data']['user_num']  # 获取用户数量self.item_num = configs['data']['item_num']  # 获取物品数量def normalizeAdj(self, adj):degree = t.pow(t.sparse.sum(adj, dim=1).to_dense() + 1e-12, -0.5)  # 计算节点度的负0.5次方newRows, newCols = adj._indices()[0, :], adj._indices()[1, :]  # 获取邻接矩阵的行和列索引rowNorm, colNorm = degree[newRows], degree[newCols]  # 获取行和列的归一化度newVals = adj._values() * rowNorm * colNorm  # 计算新的值return t.sparse.FloatTensor(adj._indices(), newVals, adj.shape)  # 返回归一化后的稀疏矩阵def forward(self, adj, seeds):rows = adj._indices()[0, :]  # 获取邻接矩阵的行索引cols = adj._indices()[1, :]  # 获取邻接矩阵的列索引maskNodes = [seeds]  # 初始化掩码节点列表for i in range(self.mask_depth):  # 遍历掩码深度curSeeds = seeds if i == 0 else nxtSeeds  # 获取当前种子节点nxtSeeds = list()  # 初始化下一个种子节点列表for seed in curSeeds:  # 遍历当前种子节点rowIdct = (rows == seed)  # 获取当前种子节点的行索引colIdct = (cols == seed)  # 获取当前种子节点的列索引idct = t.logical_or(rowIdct, colIdct)  # 合并行索引和列索引if i != self.mask_depth - 1:  # 如果不是最后一层掩码mskRows = rows[idct]  # 获取掩码后的行索引mskCols = cols[idct]  # 获取掩码后的列索引nxtSeeds.append(mskRows)  # 添加掩码后的行索引到下一个种子节点列表nxtSeeds.append(mskCols)  # 添加掩码后的列索引到下一个种子节点列表rows = rows[t.logical_not(idct)]  # 更新行索引，去掉掩码后的行cols = cols[t.logical_not(idct)]  # 更新列索引，去掉掩码后的列if len(nxtSeeds) > 0:  # 如果下一个种子节点列表不为空nxtSeeds = t.unique(t.concat(nxtSeeds))  # 合并并去重下一个种子节点列表maskNodes.append(nxtSeeds)  # 添加到掩码节点列表sampNum = int((self.user_num + self.item_num) * self.keep_rate)  # 计算采样节点数sampedNodes = t.randint(self.user_num + self.item_num, size=[sampNum]).cuda()  # 随机采样节点if self.flag == False:  # 如果标志位为False，打印信息l1 = adj._values().shape[0]  # 获取邻接矩阵的非零元素数量l2 = rows.shape[0]  # 获取掩码后的行数量print('-----')print('LENGTH CHANGE', '%.2f' % (l2 / l1), l2, l1)  # 打印长度变化tem = t.unique(t.concat(maskNodes))  # 合并并去重掩码节点列表print('Original SAMPLED NODES', '%.2f' % (tem.shape[0] / (self.user_num + self.item_num)), tem.shape[0], (self.user_num + self.item_num))  # 打印原始采样节点maskNodes.append(sampedNodes)  # 添加采样节点到掩码节点列表maskNodes = t.unique(t.concat(maskNodes))  # 合并并去重掩码节点列表if self.flag == False:  # 如果标志位为False，打印信息print('AUGMENTED SAMPLED NODES', '%.2f' % (maskNodes.shape[0] / (self.user_num + self.item_num)), maskNodes.shape[0], (self.user_num + self.item_num))  # 打印增强后的采样节点self.flag = True  # 设置标志位为Trueprint('-----')encoder_adj = self.normalizeAdj(t.sparse.FloatTensor(t.stack([rows, cols], dim=0), t.ones_like(rows).cuda(), adj.shape))  # 归一化后的编码器邻接矩阵temNum = maskNodes.shape[0]  # 获取掩码节点数量temRows = maskNodes[t.randint(temNum, size=[adj._values().shape[0]]).cuda()]  # 随机采样行索引temCols = maskNodes[t.randint(temNum, size=[adj._values().shape[0]]).cuda()]  # 随机采样列索引newRows = t.concat([temRows, temCols, t.arange(self.user_num+self.item_num).cuda(), rows])  # 合并新的行索引newCols = t.concat([temCols, temRows, t.arange(self.user_num+self.item_num).cuda(), cols])  # 合并新的列索引# 过滤重复值hashVal = newRows * (self.user_num + self.item_num) + newCols  # 计算哈希值hashVal = t.unique(hashVal)  # 去重哈希值newCols = hashVal % (self.user_num + self.item_num)  # 计算新的列索引newRows = ((hashVal - newCols) / (self.user_num + self.item_num)).long()  # 计算新的行索引decoder_adj = t.sparse.FloatTensor(t.stack([newRows, newCols], dim=0), t.ones_like(newRows).cuda().float(), adj.shape)  # 创建解码器邻接矩阵return encoder_adj, decoder_adj  # 返回编码器和解码器邻接矩阵

data_utils

data_handler_general_cf.py

import pickle  # 导入pickle模块，用于序列化和反序列化对象
import numpy as np  # 导入numpy模块，用于数值计算
from scipy.sparse import csr_matrix, coo_matrix, dok_matrix  # 从scipy.sparse中导入稀疏矩阵相关类
import scipy.sparse as sp  # 导入scipy.sparse模块，并命名为sp
from config.configurator import configs  # 从配置模块中导入配置对象configs
from data_utils.datasets_general_cf import PairwiseTrnData, AllRankTstData, PairwiseWEpochFlagTrnData  # 从数据集模块中导入相关数据集类
import torch as t  # 导入PyTorch，并命名为t
import torch.utils.data as data  # 导入PyTorch的数据工具模块class DataHandlerGeneralCF:def __init__(self):# 初始化函数，根据配置文件选择不同的数据集路径if configs['data']['name'] == 'yelp':predir = './datasets/general_cf/sparse_yelp/'elif configs['data']['name'] == 'gowalla':predir = './datasets/general_cf/sparse_gowalla/'elif configs['data']['name'] == 'amazon':predir = './datasets/general_cf/sparse_amazon/'# 定义训练、验证和测试数据文件路径self.trn_file = predir + 'train_mat.pkl'self.val_file = predir + 'valid_mat.pkl'self.tst_file = predir + 'test_mat.pkl'def _load_one_mat(self, file):"""从文件中加载一个邻接矩阵参数:file (string): 文件路径返回:scipy.sparse.coo_matrix: 加载的邻接矩阵"""with open(file, 'rb') as fs:mat = (pickle.load(fs) != 0).astype(np.float32)  # 反序列化并转换为浮点矩阵if type(mat) != coo_matrix:mat = coo_matrix(mat)  # 确保矩阵类型为coo_matrixreturn matdef _normalize_adj(self, mat):"""对邻接矩阵进行拉普拉斯归一化参数:mat (scipy.sparse.coo_matrix): 未归一化的邻接矩阵返回:scipy.sparse.coo_matrix: 归一化后的邻接矩阵"""degree = np.array(mat.sum(axis=-1)) + 1e-10  # 计算度并添加一个小值以避免除零d_inv_sqrt = np.reshape(np.power(degree, -0.5), [-1])  # 计算度的逆平方根d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.0  # 处理无穷大的值d_inv_sqrt_mat = sp.diags(d_inv_sqrt)  # 构建对角矩阵return mat.dot(d_inv_sqrt_mat).transpose().dot(d_inv_sqrt_mat).tocoo()  # 返回归一化的邻接矩阵def _make_torch_adj(self, mat):"""将单向邻接矩阵转换为双向邻接矩阵，并转换为torch.sparse.FloatTensor参数:mat (coo_matrix): 单向邻接矩阵返回:torch.sparse.FloatTensor: 双向邻接矩阵"""a = csr_matrix((configs['data']['user_num'], configs['data']['user_num']))  # 创建用户数大小的稀疏矩阵b = csr_matrix((configs['data']['item_num'], configs['data']['item_num']))  # 创建项目数大小的稀疏矩阵mat = sp.vstack([sp.hstack([a, mat]), sp.hstack([mat.transpose(), b])])  # 构建双向矩阵mat = (mat != 0) * 1.0  # 二值化mat = self._normalize_adj(mat)  # 归一化# 构建torch稀疏张量idxs = t.from_numpy(np.vstack([mat.row, mat.col]).astype(np.int64))  # 提取索引vals = t.from_numpy(mat.data.astype(np.float32))  # 提取值shape = t.Size(mat.shape)  # 获取矩阵形状return t.sparse.FloatTensor(idxs, vals, shape).to(configs['device'])  # 返回torch稀疏张量def load_data(self):# 加载训练、验证和测试数据trn_mat = self._load_one_mat(self.trn_file)tst_mat = self._load_one_mat(self.tst_file)val_mat = self._load_one_mat(self.val_file)self.trn_mat = trn_matconfigs['data']['user_num'], configs['data']['item_num'] = trn_mat.shape  # 设置用户和项目数量self.torch_adj = self._make_torch_adj(trn_mat)  # 生成torch稀疏张量# 根据训练损失类型选择训练数据集类if configs['train']['loss'] == 'pairwise':trn_data = PairwiseTrnData(trn_mat)elif configs['train']['loss'] == 'pairwise_with_epoch_flag':trn_data = PairwiseWEpochFlagTrnData(trn_mat)val_data = AllRankTstData(val_mat, trn_mat)  # 生成验证数据集tst_data = AllRankTstData(tst_mat, trn_mat)  # 生成测试数据集# 创建数据加载器self.valid_dataloader = data.DataLoader(val_data, batch_size=configs['test']['batch_size'], shuffle=False, num_workers=0)self.test_dataloader = data.DataLoader(tst_data, batch_size=configs['test']['batch_size'], shuffle=False, num_workers=0)self.train_dataloader = data.DataLoader(trn_data, batch_size=configs['train']['batch_size'], shuffle=True, num_workers=0)

输入输出说明

输入

1. 配置文件 configs 包含数据集名称、用户数、项目数、设备、批量大小和训练损失类型等配置信息。
2. 数据文件（如 train_mat.pkl, valid_mat.pkl, test_mat.pkl）包含训练、验证和测试数据的邻接矩阵。
   输出
3. torch_adj: PyTorch 的稀疏张量表示的归一化邻接矩阵，用于训练模型。
4. train_dataloader, valid_dataloader, test_dataloader: PyTorch 的数据加载器，分别用于加载训练、验证和测试数据集。

使用方法

1. 创建 DataHandlerGeneralCF 类的实例。
2. 调用 load_data 方法加载数据。
3. 使用生成的 train_dataloader, valid_dataloader, test_dataloader 进行模型训练和评估。

trainer

tuner.py

from models.bulid_model import build_model  # 导入从models.build_model模块中的build_model函数
from config.configurator import configs  # 导入config.configurator模块中的configs字典
import torch  # 导入PyTorch库
from trainer.trainer import init_seed  # 导入trainer.trainer模块中的init_seed函数class Tuner(object):def __init__(self, logger):self.logger = logger  # 初始化logger属性为提供的logger对象self.hyperparameters = configs['tune']['hyperparameters']  # 从configs字典中提取超参数self.tune_list = []  # 初始化一个空列表，用于存储超参数的值self.search_length = 1  # 初始化搜索长度为1# 遍历每个超参数，并初始化tune_list和search_lengthfor hyper_para in self.hyperparameters:self.tune_list.append(configs['tune'][hyper_para])  # 将超参数的值添加到tune_list中self.search_length = self.search_length * len(configs['tune'][hyper_para])  # 计算总搜索长度# 计算参数长度并初始化hex_length用于索引self.para_length = [len(para_list) for para_list in self.tune_list]self.hex_length = [1 for _ in range(len(self.tune_list))]for i in range(len(self.para_length) - 2, -1, -1):self.hex_length[i] = self.para_length[i + 1] * self.hex_length[i + 1]self.origin_model_para = configs['model'].copy()  # 复制原始模型参数def zero_step(self):self.now_step = 0  # 将now_step属性初始化为0def step(self):self.now_step += 1  # 将now_step属性增加1def next_model(self, data_handler):init_seed()  # 初始化种子以保证可重复性now_para = {}  # 初始化一个空字典用于存储当前超参数now_para_str = ''  # 初始化一个空字符串用于存储连接的超参数名和值# 遍历每个超参数for i in range(len(self.hyperparameters)):para_name = self.hyperparameters[i]  # 获取当前超参数的名称selected_idx = (self.now_step // self.hex_length[i]) % self.para_length[i]  # 计算当前超参数值的索引selected_val = self.tune_list[i][selected_idx]  # 获取当前超参数的选定值now_para[para_name] = selected_val  # 在now_para字典中存储当前超参数名和值now_para_str += '{}{}'.format(para_name, selected_val)  # 将当前超参数名和值连接到now_para_str中configs['model'][para_name] = selected_val  # 使用当前超参数值更新模型配置configs['tune']['now_para_str'] = now_para_str  # 更新configs，存储连接的当前超参数名和值的字符串self.logger.log('hyperparameter: {}'.format(now_para))  # 记录当前超参数设置model = build_model(data_handler).cuda()  # 使用build_model函数构建模型并移动到CUDAreturn model  # 返回构建的模型def grid_search(self, data_handler, trainer):self.zero_step()  # 将搜索步骤初始化为零# 遍历所有超参数值的组合for _ in range(self.search_length):model = self.next_model(data_handler)  # 获取下一个模型配置trainer.train(model)  # 使用提供的训练器训练模型# trainer.evaluate(model)  # 可选：评估模型性能del model  # 删除模型以释放GPU内存torch.cuda.empty_cache()  # 清空CUDA内存缓存self.step()  # 进入网格搜索的下一个步骤configs['model'] = self.origin_model_para.copy()  # 在网格搜索后恢复原始模型参数

这段代码定义了一个Tuner类，用于使用网格搜索方法进行超参数调优。它遍历超参数值的组合，根据配置构建和训练模型，并管理超参数的配置和记录。