SASRec论文代码复现讲解（教你如何读代码）倾囊相授！

坚持每天敲代码

已于 2024-04-25 17:08:16 修改

阅读量2.5k

点赞数 19

文章标签：推荐算法论文笔记人工智能 pytorch

于 2024-04-22 15:08:24 首次发布

本文链接：https://blog.youkuaiyun.com/m0_59772307/article/details/137956994

版权

SASRec

paper: Self-Attentive Sequential Recommendation

publication：2018 IEEE International Conference on Data Mining

本文是我个人的学习总结，我也是初学者，但是我有很认真在思考做笔记，如果你也是初学者，我相信这对你的帮助很大，因为下面是作为刚入门的我的思考过程，这次先讲代码中部分函数，后边我会在本文慢慢更新，记得关注。可以先试着看看下边的内容，或许对你有所帮助！在此之前，我希望你能去仔细读一下原文，带着你疑惑的地方来看看代码如何实现，我强烈建议刚入门的学者可以手动一行一行的敲一下，这真的很管用！话不多说，上代码讲解！ transformer的文章可以看我之前发的一篇：Transformer代码（Pytorch实现和详解）！！！-优快云博客

我想讲一下我读代码的方法：目前，我比较喜欢一个模块一个模块 or 一个函数一个函数去解决，我可能先看数据怎么处理划分的，每个模块我会用代码测试一下，可以看一下输出结果跟自己想的一不一样，遇到不会的函数就去百度+自己写代码测试，非常好用。然后看懂一个模块之后，自己手动敲一遍，这看起来比较费事，但是对你的代码能力提升非常快！

DataSet

MovieLens-1M 是电影评分数据集，地址：MovieLens 1M Dataset | GroupLens

但是作者应该是把数据集处理成了下边的样子，形式就是(user, item)

user-用户ID；item-电影ID

我去官网看了一下，应该都是5星级的评分，所以这里的数据都表示用户喜欢的电影，我猜的~

非常简单对不对

Arguments

备注：这里如果不知道argparse的话，可以去了解一下。


parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=False,default = 'ml-1m')
parser.add_argument('--train_dir', required=False,default = 'default')
parser.add_argument('--batch_size', default=128, type=int)
parser.add_argument('--lr', default=0.001, type=float)
parser.add_argument('--maxlen', default=50, type=int)
parser.add_argument('--hidden_units', default=50, type=int)
parser.add_argument('--num_blocks', default=2, type=int)
parser.add_argument('--num_epochs', default=201, type=int)
parser.add_argument('--num_heads', default=1, type=int)
parser.add_argument('--dropout_rate', default=0.5, type=float)
parser.add_argument('--l2_emb', default=0.0, type=float)
parser.add_argument('--device', default='cpu', type=str)
parser.add_argument('--inference_only', default=False, type=str2bool)
parser.add_argument('--state_dict_path', default=None, type=str)

args = parser.parse_args()

我们通过args来管理我们模型用到的参数，如果是需要训练之前手动设置的参数叫“超参数”，还有一种是模型训练过程中得到的参数。这里就是超参数：

dataset：数据集名称，默认是“ml-1m”

train_dir：训练目录，默认是“default”，后边路径会拼接

batch_size：批量大小，就是一个批量有多少个样本。

lr：学习率

maxlen：用户序列长度（user:[item1, item2, item3…………] 就是item的长度）

hidden_units：隐藏层，就是项目和位置嵌入的维度。

num_blocks：transformer中提到

num_epochs：训练轮数。一轮的话就是需要将所有的样本过一遍。

num_heads：注意力机制的头数

dropout_rate：dropout率

l2_emb：这个暂时不知道

device：设备，是“cpu”或者“cuda”，cuda就是GPU

inference_only：这里inference指的是用训练好的模型预测，不需要再进行训练。

state_dict_path：保存模型之后的state_dict的文件路径，用来加载训练好的模型。

data_partition

这个函数的主要作用是将数据集划分成训练集、验证集、测试集，仔细看看怎么操作的。

from collections import defaultdict

def data_partition(fname):
    usernum = 0
    itemnum = 0
    User = defaultdict(list)
    #默认字典,key不存在时返回默认值。 {user:[itme1,item2……]}
    user_train = {}
    user_valid = {}
    user_test = {}
    # assume user/item index starting from 1
    f = open('data/%s.txt' % fname, 'r')
    for line in f:
        u, i = line.rstrip().split(' ')
        u = int(u)
        i = int(i)
        #usernum、itemnum记录最大的用户、项目编号，正好是用户、项目的总数量，因为，都是从1开始编号的。
        usernum = max(u, usernum)
        itemnum = max(i, itemnum)
        User[u].append(i)

    for user in User:
        #当前用户的序列长度
        nfeedback = len(User[user])
        #如果当前用户序列长度小于3
        if nfeedback < 3:
            #把当前user的交互序列都给了训练集，验证和测试集为空
            user_train[user] = User[user]
            user_valid[user] = []
            user_test[user] = []
        else:
            user_train[user] = User[user][:-2]
            #让字典的value值是列表，而不是单个元素
            user_valid[user] = []
            user_valid[user].append(User[user][-2])
            user_test[user] = []
            user_test[user].append(User[user][-1])
    return [user_train, user_valid, user_test, usernum, itemnum]


#test
user_train, user_valid, user_test, usernum, itemnum = data_partition('ml-1m')
print('========================user_valid================================')
print(user_valid)
print('========================user_test================================')
print(user_test)

文件路径，以及该如何读文件这都很常规的操作了吧。
去了解并测试不懂的函数，比如：defaultdict()、rstrip()。
代码的逻辑很简单，主要是各个数据的形式。你可以打印去看一下User、user_train、user_valid、user_test，记住他们的样子。

测试技巧

如果你对于测试无从下手，可以试试以下的步骤：

可以试着先创建一个单独的test.py文件,用来测试。
将代码粘贴过来，包括需要的库，以及准备好数据集，当然你也可以自己对代码稍作改动。
如果你想知道中间某个变量的值，可以使用print函数打印出来；针对for循环里边的语句，可以打断点，然后调试，一层一层的看变量值。
针对不懂的函数，比如defaultdict(), 先去百度了解功能和使用，再去单独测试，再回去看代码！

我想这对你来说小菜一叠~~

知识补充！

正样本：用户交互过的项目

负样本：用户还没有交互过的项目

WarpSampler

这一部分主要是进行样本采样，会结合队列Queue和多进程来实现，大致就是创建多个进程，每个进程都会采集一个batch个用户样本（包括用户的正负样本），放到Queue中，如果Queue满了，采样的进程就会先阻塞。WarpSampler中的函数next_batch每次能取一个batch的样本数据。

#随机采样一个负样本
def random_neq(l, r, s):
    t = np.random.randint(l, r)
    while t in s:
        t = np.random.randint(l, r)
    return t

#采样函数
def sample_function(user_train, usernum, itemnum, batch_size, maxlen, result_queue, SEED):
    def sample():
        '''
        :return: (user, seq, pos, neg)
        '''
        user = np.random.randint(1, usernum+1)
        while len(user_train[user]) <= 1 : user = np.random.randint(1, usernum+1)

        seq = np.zeros([maxlen], dtype = np.int32)
        pos = np.zeros([maxlen], dtype = np.int32)
        neg = np.zeros([maxlen], dtype = np.int32)
        nxt = user_train[user][-1]
        idx = maxlen - 1

        ts = set(user_train[user]) # 当前user序列中交互过的item集合
        for i in reversed(user_train[1][:-1]):
            seq[idx] = i
            pos[idx] = nxt
            if nxt != 0: neg[idx] = random_neq(1, itemnum + 1, ts)
            nxt = i
            idx -= 1
            if idx == -1: break

        return (user, seq, pos, neg)

    np.random.seed(SEED)
    while True:
        one_batch = []
        print("##############")
        for i in range(batch_size):
            one_batch.append(sample())
        #queue满了会阻塞
        result_queue.put(zip(*one_batch))

class WarpSampler(object):
    def __init__(self, User, usernum, itemnum, batch_size=64, maxlen=10, n_workers=1):
        self.result_queue = Queue(maxsize = n_workers*10)
        self.processors = []
        for i in range(n_workers):
            self.processors.append(
                Process(target = sample_function, args=(User,
                                                        usernum,
                                                        itemnum,
                                                        batch_size,
                                                        maxlen,
                                                        self.result_queue,
                                                        np.random.randint(2e9)
                                                        )))
            self.processors[-1].daemon = True
            self.processors[-1].start()
    def next_batch(self):
        return self.result_queue.get()

    def close(self):
        for p in self.processors:
            p.terminate()
            p.join()

SASRec

拿到模型主干之后，先看 init() 里边都定义了啥。这里先定义了usernum、itemnum、device。接着是两个嵌入层：item_emb 和 pos_emb，他们的嵌入维度都是hidden_units，还有一个dropout层，用来对嵌入之后的结果随机失活。定义了四个ModuleList来存储各个层。num_blocks是堆叠层的数量，每一层有两个子层：一个多头自注意力层，后跟LayerNorm；一个前馈网络层，后跟LayerNorm。

第二步看一下forward() 里边如何组织网络的，这里直接调用log2feats方法，我们进到该方法里边看。里边的操作都是transformer的正常操作，从上到下一次是：

这里输入序列log_seq的样子：{ user1 : [item1, item2, item3……] }

 def log2feats(self, log_seqs):
    seqs = self.item_emb(torch.LongTensor(log_seqs).to(self.dev))#项目嵌入
    seqs *= self.item_emb.embedding_dim ** 0.5#缩放
    positions = np.tile(np.array(range(log_seqs.shape[1])), [log_seqs.shape[0], 1])#生成位置编码
    seqs += self.pos_emb(torch.LongTensor(positions).to(self.dev))#将位置编码进行位置嵌入，并且加到项目的嵌入向量序列中，位置感知
    seqs = self.emb_dropout(seqs)#dropout操作

如上代码所示，首先对输入seq进行item_emb，然后对seq进行缩放；自定义生成位置编码，对位置编码进行pos_emb, 然后将位置嵌入和项目嵌入想加，这样将位置信息也融合到序列里边了。对嵌入之后的序列进行dropout操作。这都是常规操作了吧。

timeline_mask = torch.BoolTensor(log_seqs == 0).to(self.dev)         

seqs *= ~timeline_mask.unsqueeze(-1)

上面这俩行代码的意思是，首先创建一个布尔张量，形状和log_seqs一样，只不过log_seqs中等于0的位置置为True，其他为False。将掩码在最后一个维度进行扩展以便与seqs相乘进行广播操作，并将他取反，也就是True表示log_seq不等于0的位置。综上所述就是将嵌入之后的seqs的某些位置置为0了，这些位置正好是原序列log_seq中item等于0的位置。

tl = seqs.shape[1] # time dim len for enforce causality
attention_mask = ~torch.tril(torch.ones((tl, tl), dtype=torch.bool, device=self.dev))

上面代码，创建注意力掩码矩阵，transformer中有讲，它创建了一个下三角形矩阵，将下三角部分置为 True，其余部分置为 False，并且使用取反操作符 ~ 将其取反，以便在注意力计算中将这些位置的值屏蔽掉。tl取的是seq的第二个dim的长度，也就是用户序列中项目的长度，表示每个项目对其他项目的注意力是多少。为什么是三角形状呢，因为预测当前位置的时候要屏蔽掉之后的信息。

        for i in range(len(self.attention_layers)):
            seqs = torch.transpose(seqs, 0, 1)#转置操作
            Q = self.attention_layernorms[i](seqs)#将seq传入注意力层的归一化层，输出q值
            mha_outputs, _ = self.attention_layers[i](Q, seqs, seqs, #多头注意力输出
                                            attn_mask=attention_mask)#注意力掩码
                                            # key_padding_mask=timeline_mask
                                            # need_weights=False) this arg do not work?
            seqs = Q + mha_outputs#残差连接
            seqs = torch.transpose(seqs, 0, 1)

            seqs = self.forward_layernorms[i](seqs)
            seqs = self.forward_layers[i](seqs)
            seqs *=  ~timeline_mask.unsqueeze(-1)

        log_feats = self.last_layernorm(seqs) # (U, T, C) -> (U, -1, C)

这段代码遍历num_blocks个encoder层，每个encoder层，先后分别进行了attention的层归一化，self-attention层, 参差连接，feedforward的层归一化，feedforward层。最后对输出结果在进行一次层归一化。输出结果。

multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
attn_output, attn_output_weights = multihead_attn(query, key, value)

这个是官网调用多头注意力返回的两个输出结果，一个是注意力输出，一个是注意力权重，这里我们没有用到权重就直接忽略了。

至此，我们得到了经过Encoder之后的输出结果。

    def forward(self, user_ids, log_seqs, pos_seqs, neg_seqs): # for training        
        log_feats = self.log2feats(log_seqs) # user_ids hasn't been used yet

        pos_embs = self.item_emb(torch.LongTensor(pos_seqs).to(self.dev))
        neg_embs = self.item_emb(torch.LongTensor(neg_seqs).to(self.dev))

        pos_logits = (log_feats * pos_embs).sum(dim=-1)
        neg_logits = (log_feats * neg_embs).sum(dim=-1)

        # pos_pred = self.pos_sigmoid(pos_logits)
        # neg_pred = self.neg_sigmoid(neg_logits)

        return pos_logits, neg_logits # pos_pred, neg_pred

这是forward方法，首先是让原始的用户-项目序列经过了自注意力得到特征表示（QKV一系列操作）：log_feats, 然后对正负样本进行嵌入，将log_feats和pos_emb、neg_emb分别逐点相乘，在最后一个维度求和，得到正负样本的逻辑回归得分。

    def predict(self, user_ids, log_seqs, item_indices): # for inference
        log_feats = self.log2feats(log_seqs) # user_ids hasn't been used yet

        final_feat = log_feats[:, -1, :] # only use last QKV classifier, a waste

        item_embs = self.item_emb(torch.LongTensor(item_indices).to(self.dev)) # (U, I, C)

        logits = item_embs.matmul(final_feat.unsqueeze(-1)).squeeze(-1)

        # preds = self.pos_sigmoid(logits) # rank same item list for different users

        return logits # preds # (U, I)

同样还是将用户序列，经过自注意力转换为特征表示。final_feat = log_feats[:, -1, :] 这里是取最后一个时间步，[batch_size, seq_len, d_model],这里就是指的seq_len的最后一个item。后边的可以打断点看看里边是什么样子。

SASRec完整代码：

class SASRec(torch.nn.Module):
    def __init__(self, user_num, item_num, args):
        super(SASRec, self).__init__()

        self.user_num = user_num # 用户数量
        self.item_num = item_num # 项目数量
        self.dev = args.device   # 设备参数

        # TODO: loss += args.l2_emb for regularizing embedding vectors during training
        # https://stackoverflow.com/questions/42704283/adding-l1-l2-regularization-in-pytorch
        #嵌入层
        self.item_emb = torch.nn.Embedding(self.item_num+1, args.hidden_units, padding_idx=0)
        self.pos_emb = torch.nn.Embedding(args.maxlen, args.hidden_units) # TO IMPROVE
        #定义Dropout层
        self.emb_dropout = torch.nn.Dropout( p = args.dropout_rate)

        #ModuleList
        self.attention_layernorms = torch.nn.ModuleList()
        self.attention_layers = torch.nn.ModuleList()
        self.forward_layernorms = torch.nn.ModuleList()
        self.forward_layers = torch.nn.ModuleList()

        self.last_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)

        for _ in range(args.num_blocks):
            #注意力-层归一化
            new_attn_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)
            self.attention_layernorms.append(new_attn_layernorm)

            #多头注意力层
            new_attn_layer = torch.nn.MultiheadAttention(args.hidden_units,
                                                            args.num_heads,
                                                            args.dropout_rate)
            self.attention_layers.append(new_attn_layer)

            #FeedForward-层归一化
            new_fwd_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)
            self.forward_layernorms.append(new_fwd_layernorm)

            #FeedForward层
            new_fwd_layer = PointWiseFeedForward(args.hidden_units, args.dropout_rate)
            self.forward_layers.append(new_fwd_layer)

            # self.pos_sigmoid = torch.nn.Sigmoid()
            # self.neg_sigmoid = torch.nn.Sigmoid()

    
    def log2feats(self, log_seqs):
        seqs = self.item_emb(torch.LongTensor(log_seqs).to(self.dev))#项目嵌入
        seqs *= self.item_emb.embedding_dim ** 0.5#缩放
        positions = np.tile(np.array(range(log_seqs.shape[1])), [log_seqs.shape[0], 1])#生成位置编码
        seqs += self.pos_emb(torch.LongTensor(positions).to(self.dev))#将位置编码进行位置嵌入，并且加到项目的嵌入向量序列中，位置感知
        seqs = self.emb_dropout(seqs)#dropout操作

        timeline_mask = torch.BoolTensor(log_seqs == 0).to(self.dev)#
        seqs *= ~timeline_mask.unsqueeze(-1) # broadcast in last dim

        tl = seqs.shape[1] # time dim len for enforce causality
        attention_mask = ~torch.tril(torch.ones((tl, tl), dtype=torch.bool, device=self.dev))

        for i in range(len(self.attention_layers)):
            seqs = torch.transpose(seqs, 0, 1)#转置操作
            Q = self.attention_layernorms[i](seqs)#将seq传入注意力层的归一化层，输出q值
            mha_outputs, _ = self.attention_layers[i](Q, seqs, seqs, #多头注意力输出
                                            attn_mask=attention_mask)#注意力掩码
                                            # key_padding_mask=timeline_mask
                                            # need_weights=False) this arg do not work?
            seqs = Q + mha_outputs#残差连接
            seqs = torch.transpose(seqs, 0, 1)

            seqs = self.forward_layernorms[i](seqs)
            seqs = self.forward_layers[i](seqs)
            seqs *=  ~timeline_mask.unsqueeze(-1)

        log_feats = self.last_layernorm(seqs) # (U, T, C) -> (U, -1, C)

        return log_feats
    #前向传播逻辑
    def forward(self, user_ids, log_seqs, pos_seqs, neg_seqs): # for training        
        log_feats = self.log2feats(log_seqs) # user_ids hasn't been used yet

        pos_embs = self.item_emb(torch.LongTensor(pos_seqs).to(self.dev))
        neg_embs = self.item_emb(torch.LongTensor(neg_seqs).to(self.dev))

        pos_logits = (log_feats * pos_embs).sum(dim=-1)
        neg_logits = (log_feats * neg_embs).sum(dim=-1)

        # pos_pred = self.pos_sigmoid(pos_logits)
        # neg_pred = self.neg_sigmoid(neg_logits)

        return pos_logits, neg_logits # pos_pred, neg_pred

    def predict(self, user_ids, log_seqs, item_indices): # for inference
        log_feats = self.log2feats(log_seqs) # user_ids hasn't been used yet

        final_feat = log_feats[:, -1, :] # only use last QKV classifier, a waste

        item_embs = self.item_emb(torch.LongTensor(item_indices).to(self.dev)) # (U, I, C)

        logits = item_embs.matmul(final_feat.unsqueeze(-1)).squeeze(-1)

        # preds = self.pos_sigmoid(logits) # rank same item list for different users

        return logits # preds # (U, I)