推荐系统（八）—— 知识图谱增强推荐系统

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

import torch
import torch.nn as nn
from collections import defaultdict
from typing import List, Tuple


class KnowledgeGraph:
    """用于推荐系统的简易知识图谱。"""

    def __init__(self):
        self.entities = {}                   # entity_name -> entity_id
        self.relations = {}                  # relation_name -> relation_id
        self.triples = []                    # [(h_id, r_id, t_id), ...]
        self.entity_adj = defaultdict(list)  # entity_id -> [(r_id, t_id), ...]

    def add_entity(self, name: str) -> int:
        if name not in self.entities:
            self.entities[name] = len(self.entities)
        return self.entities[name]

    def add_relation(self, name: str) -> int:
        if name not in self.relations:
            self.relations[name] = len(self.relations)
        return self.relations[name]

    def add_triple(self, head: str, relation: str, tail: str):
        h_id = self.add_entity(head)
        r_id = self.add_relation(relation)
        t_id = self.add_entity(tail)
        self.triples.append((h_id, r_id, t_id))
        self.entity_adj[h_id].append((r_id, t_id))

    def get_neighbors(self, entity_id: int) -> List[Tuple[int, int]]:
        return self.entity_adj[entity_id]


# 构建一个小型电影知识图谱
kg = KnowledgeGraph()
kg.add_triple("黑暗骑士", "directed_by", "诺兰")
kg.add_triple("黑暗骑士", "starred", "贝尔")
kg.add_triple("黑暗骑士", "has_genre", "动作")
kg.add_triple("黑暗骑士", "has_genre", "犯罪")
kg.add_triple("盗梦空间", "directed_by", "诺兰")
kg.add_triple("盗梦空间", "starred", "迪卡普里奥")
kg.add_triple("盗梦空间", "has_genre", "科幻")
kg.add_triple("致命魔术", "directed_by", "诺兰")
kg.add_triple("致命魔术", "starred", "贝尔")
kg.add_triple("致命魔术", "has_genre", "剧情")

# 查看图谱
dk_id = kg.entities["黑暗骑士"]
print(f"黑暗骑士的邻居: {kg.get_neighbors(dk_id)}")
# (directed_by, 诺兰), (starred, 贝尔), (has_genre, 动作), (has_genre, 犯罪)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

import torch
import torch.nn as nn
import torch.nn.functional as F


class RippleNet(nn.Module):
    """RippleNet：在知识图谱上传播用户偏好。"""

    def __init__(self, num_users, num_items, num_entities, num_relations,
                 embedding_dim=64, n_hop=2, n_memory=32):
        super().__init__()
        self.n_hop = n_hop
        self.n_memory = n_memory
        self.embedding_dim = embedding_dim

        self.user_emb = nn.Embedding(num_users, embedding_dim)
        self.item_emb = nn.Embedding(num_items, embedding_dim)
        self.entity_emb = nn.Embedding(num_entities, embedding_dim)
        self.relation_emb = nn.Embedding(num_relations, embedding_dim)

        self.hop_weights = nn.Parameter(torch.ones(n_hop + 1) / (n_hop + 1))

        for emb in [self.user_emb, self.item_emb,
                    self.entity_emb, self.relation_emb]:
            nn.init.xavier_uniform_(emb.weight)

    def forward(self, user_ids, item_ids, ripple_sets):
        """
        Args:
            user_ids: [batch_size]
            item_ids: [batch_size]
            ripple_sets: ripple_sets[i][h] 是用户 i 在第 h 跳的
                         (relation_id, tail_id) 列表。
        """
        user_base = self.user_emb(user_ids)        # [B, d]
        item_vec = self.item_emb(item_ids)          # [B, d]

        user_enhanced = self._propagate(user_base, item_vec, ripple_sets)
        scores = (user_enhanced * item_vec).sum(dim=1)
        return scores

    def _propagate(self, user_base, item_vec, ripple_sets):
        batch_size = user_base.size(0)
        hop_memories = [user_base]

        for hop in range(self.n_hop):
            hop_embs = []
            for i in range(batch_size):
                rs = ripple_sets[i][hop] if hop < len(ripple_sets[i]) else []
                if not rs:
                    hop_embs.append(torch.zeros(self.embedding_dim))
                    continue

                rs = rs[:self.n_memory]                  # 限制 memory
                rels = torch.LongTensor([r for r, _ in rs])
                tails = torch.LongTensor([t for _, t in rs])

                rel_e = self.relation_emb(rels)          # [K, d]
                tail_e = self.entity_emb(tails)          # [K, d]

                # 用候选物品对每个 KG 邻居打分
                scores = (item_vec[i].unsqueeze(0) * rel_e * tail_e).sum(1)
                probs = F.softmax(scores, dim=0)         # [K]
                hop_embs.append((probs.unsqueeze(1) * tail_e).sum(0))

            hop_memories.append(torch.stack(hop_embs))

        # 多跳加权融合
        stacked = torch.stack(hop_memories, dim=1)       # [B, H+1, d]
        weights = F.softmax(self.hop_weights, dim=0)
        return (weights.unsqueeze(0).unsqueeze(2) * stacked).sum(dim=1)

    @staticmethod
    def build_ripple_sets(user_items, kg_adj, max_hops=2):
        """通过 BFS 为单个用户构建 ripple sets。"""
        ripple_sets = []
        current = set(user_items)
        for _ in range(max_hops):
            next_hop = []
            for entity in current:
                for r_id, t_id in kg_adj.get(entity, []):
                    next_hop.append((r_id, t_id))
            ripple_sets.append(next_hop)
            current = {t for _, t in next_hop}
        return ripple_sets

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

import torch
import torch.nn as nn
import torch.nn.functional as F


class KGCNLayer(nn.Module):
    """单层 KGCN，带关系感知聚合。"""

    def __init__(self, embedding_dim, num_relations):
        super().__init__()
        self.embedding_dim = embedding_dim

        self.rel_transforms = nn.ModuleList([
            nn.Linear(embedding_dim, embedding_dim, bias=False)
            for _ in range(num_relations)
        ])
        self.attn = nn.ModuleList([
            nn.Linear(embedding_dim, 1)
            for _ in range(num_relations)
        ])
        self.combine = nn.Linear(embedding_dim * 2, embedding_dim)

    def forward(self, item_embs, neighbors, relation_embs):
        num_items = item_embs.size(0)
        aggregated = []

        for i in range(num_items):
            nbrs = neighbors.get(i, [])
            if not nbrs:
                aggregated.append(torch.zeros(self.embedding_dim))
                continue

            nbr_embs, scores = [], []
            for r_id, n_id in nbrs:
                transformed = self.rel_transforms[r_id](item_embs[n_id])
                nbr_embs.append(transformed)
                compatibility = item_embs[i] + relation_embs[r_id]
                scores.append(self.attn[r_id](compatibility))

            nbr_embs = torch.stack(nbr_embs)
            scores = torch.cat(scores)
            weights = F.softmax(scores, dim=0)
            aggregated.append((weights.unsqueeze(1) * nbr_embs).sum(0))

        aggregated = torch.stack(aggregated)
        combined = torch.cat([item_embs, aggregated], dim=1)
        return self.combine(combined)


class KGCN(nn.Module):
    """用于推荐的知识图谱卷积网络。"""

    def __init__(self, num_users, num_items, num_entities, num_relations,
                 embedding_dim=64, n_layers=2):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items

        self.user_emb = nn.Embedding(num_users, embedding_dim)
        self.entity_emb = nn.Embedding(num_entities, embedding_dim)
        self.relation_emb = nn.Embedding(num_relations, embedding_dim)

        self.layers = nn.ModuleList([
            KGCNLayer(embedding_dim, num_relations)
            for _ in range(n_layers)
        ])

        for emb in [self.user_emb, self.entity_emb, self.relation_emb]:
            nn.init.xavier_uniform_(emb.weight)

    def forward(self, user_ids, item_ids, item_neighbors):
        user_vec = self.user_emb(user_ids)
        item_vec = self.entity_emb.weight[:self.num_items]
        rel_vec = self.relation_emb.weight

        for idx, layer in enumerate(self.layers):
            if idx < len(item_neighbors):
                item_vec = layer(item_vec, item_neighbors[idx], rel_vec)
                item_vec = F.relu(item_vec)

        item_final = item_vec[item_ids]
        return (user_vec * item_final).sum(dim=1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

import torch
import torch.nn as nn
import torch.nn.functional as F


class KGATLayer(nn.Module):
    """单层 KGAT，带关系感知注意力。"""

    def __init__(self, embedding_dim, num_heads=1, dropout=0.1):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads
        self.dropout = dropout

        # 在 [entity || relation || neighbor] 上打分
        self.attention = nn.Linear(embedding_dim * 3, num_heads)
        self.W_v = nn.Linear(embedding_dim, embedding_dim)
        self.output = nn.Linear(embedding_dim, embedding_dim)

    def forward(self, entity_embs, relation_embs, neighbors):
        outputs = []
        for i in range(entity_embs.size(0)):
            nbrs = neighbors.get(i, [])
            if not nbrs:
                outputs.append(entity_embs[i])
                continue

            scores, values = [], []
            for r_id, n_id in nbrs:
                combined = torch.cat([
                    entity_embs[i], relation_embs[r_id], entity_embs[n_id]
                ])
                scores.append(self.attention(combined))
                values.append(self.W_v(entity_embs[n_id]))

            scores = F.softmax(torch.stack(scores), dim=0)  # [K, heads]
            scores = F.dropout(scores, p=self.dropout, training=self.training)
            values = torch.stack(values)                     # [K, d]

            aggregated = (scores.mean(dim=1, keepdim=True) * values).sum(0)
            out = entity_embs[i] + aggregated
            outputs.append(F.relu(self.output(out)))

        return torch.stack(outputs)


class KGAT(nn.Module):
    """用于推荐的知识图谱注意力网络。"""

    def __init__(self, num_users, num_items, num_entities, num_relations,
                 embedding_dim=64, n_layers=2, num_heads=2, dropout=0.1):
        super().__init__()
        self.num_users = num_users

        self.entity_emb = nn.Embedding(num_entities, embedding_dim)
        self.user_emb = nn.Embedding(num_users, embedding_dim)
        self.relation_emb = nn.Embedding(num_relations, embedding_dim)

        self.layers = nn.ModuleList([
            KGATLayer(embedding_dim, num_heads, dropout)
            for _ in range(n_layers)
        ])

        for emb in [self.entity_emb, self.user_emb, self.relation_emb]:
            nn.init.xavier_uniform_(emb.weight)

    def forward(self, user_ids, item_ids, ckg_neighbors):
        # 合并嵌入表：[users | entities]
        all_embs = self.entity_emb.weight.clone()
        all_embs[:self.num_users] = self.user_emb.weight
        rel_embs = self.relation_emb.weight

        for idx, layer in enumerate(self.layers):
            if idx < len(ckg_neighbors):
                all_embs = layer(all_embs, rel_embs, ckg_neighbors[idx])

        user_final = all_embs[user_ids]
        item_final = all_embs[self.num_users + item_ids]
        return (user_final * item_final).sum(dim=1)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

import torch
import torch.nn as nn
import torch.nn.functional as F


class TransR(nn.Module):
    """TransR：每个关系自带一套实体投影矩阵。"""

    def __init__(self, num_entities, num_relations, entity_dim, relation_dim):
        super().__init__()
        self.entity_emb = nn.Embedding(num_entities, entity_dim)
        self.relation_emb = nn.Embedding(num_relations, relation_dim)
        self.proj_matrices = nn.Embedding(num_relations,
                                          entity_dim * relation_dim)
        self.entity_dim = entity_dim
        self.relation_dim = relation_dim

        for emb in [self.entity_emb, self.relation_emb, self.proj_matrices]:
            nn.init.xavier_uniform_(emb.weight)

    def forward(self, heads, relations, tails,
                neg_heads=None, neg_tails=None):
        h = self.entity_emb(heads)
        r = self.relation_emb(relations)
        t = self.entity_emb(tails)

        proj = self.proj_matrices(relations).view(
            -1, self.entity_dim, self.relation_dim
        )
        h_proj = torch.bmm(h.unsqueeze(1), proj).squeeze(1)
        t_proj = torch.bmm(t.unsqueeze(1), proj).squeeze(1)

        pos_score = (h_proj + r - t_proj).norm(p=2, dim=1) ** 2

        if neg_heads is not None and neg_tails is not None:
            nh = self.entity_emb(neg_heads)
            nt = self.entity_emb(neg_tails)
            nh_proj = torch.bmm(nh.unsqueeze(1), proj).squeeze(1)
            nt_proj = torch.bmm(nt.unsqueeze(1), proj).squeeze(1)
            neg_score = (nh_proj + r - nt_proj).norm(p=2, dim=1) ** 2
            return F.relu(1.0 + pos_score - neg_score).mean()

        return pos_score


class CKE(nn.Module):
    """协同知识库嵌入。"""

    def __init__(self, num_users, num_items, num_entities, num_relations,
                 embedding_dim=64, relation_dim=32,
                 kg_lambda=0.1, reg_lambda=0.01):
        super().__init__()
        self.num_entities = num_entities
        self.kg_lambda = kg_lambda
        self.reg_lambda = reg_lambda

        # CF 分量
        self.user_emb = nn.Embedding(num_users, embedding_dim)
        self.item_cf_emb = nn.Embedding(num_items, embedding_dim)

        # KG 分量（TransR）
        self.transr = TransR(num_entities, num_relations,
                             embedding_dim, relation_dim)

        # 把 KG 嵌入投影到 CF 空间
        self.kg_proj = nn.Linear(embedding_dim, embedding_dim)

        nn.init.xavier_uniform_(self.user_emb.weight)
        nn.init.xavier_uniform_(self.item_cf_emb.weight)

    def forward(self, user_ids, item_ids):
        user_vec = self.user_emb(user_ids)
        item_cf = self.item_cf_emb(item_ids)
        item_kg = self.kg_proj(self.transr.entity_emb(item_ids))

        item_combined = item_cf + item_kg
        return (user_vec * item_combined).sum(dim=1)

    def compute_loss(self, user_ids, item_ids, ratings,
                     kg_heads=None, kg_rels=None, kg_tails=None):
        # CF 损失
        preds = self.forward(user_ids, item_ids)
        cf_loss = F.mse_loss(preds, ratings.float())

        # KG 损失
        kg_loss = torch.tensor(0.0)
        if kg_heads is not None:
            neg_heads = torch.randint(0, self.num_entities, kg_heads.shape)
            neg_tails = torch.randint(0, self.num_entities, kg_tails.shape)
            kg_loss = self.transr(kg_heads, kg_rels, kg_tails,
                                  neg_heads, neg_tails)

        # 正则
        reg = (self.user_emb.weight.norm(2) ** 2 +
               self.item_cf_emb.weight.norm(2) ** 2)

        return cf_loss + self.kg_lambda * kg_loss + self.reg_lambda * reg

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import deque


class PathReasoning(nn.Module):
    """寻找并打分知识图谱路径，用于可解释推荐。"""

    def __init__(self, num_entities, num_relations, embedding_dim=64):
        super().__init__()
        self.entity_emb = nn.Embedding(num_entities, embedding_dim)
        self.relation_emb = nn.Embedding(num_relations, embedding_dim)

        self.path_scorer = nn.Sequential(
            nn.Linear(embedding_dim * 3, embedding_dim),
            nn.ReLU(),
            nn.Linear(embedding_dim, 1),
        )

        nn.init.xavier_uniform_(self.entity_emb.weight)
        nn.init.xavier_uniform_(self.relation_emb.weight)

    def find_paths(self, kg_adj, start, end, max_hops=3, max_paths=10):
        """BFS 查找从 start 到 end 的路径。"""
        paths = []
        queue = deque([(start, [], [])])
        visited = set()

        while queue and len(paths) < max_paths:
            current, rel_path, ent_path = queue.popleft()

            if current == end and rel_path:
                paths.append((rel_path, ent_path + [end]))
                continue
            if len(rel_path) >= max_hops:
                continue

            for r_id, t_id in kg_adj.get(current, []):
                if (current, r_id, t_id) not in visited:
                    visited.add((current, r_id, t_id))
                    queue.append((t_id,
                                  rel_path + [r_id],
                                  ent_path + [current]))

        return paths[:max_paths]

    def score_path(self, path):
        """用嵌入对单条路径打分。"""
        rel_ids, ent_ids = path
        if not rel_ids:
            return torch.tensor(0.0)

        start = self.entity_emb(torch.tensor([ent_ids[0]]))
        rels = self.relation_emb(torch.tensor(rel_ids)).mean(0, keepdim=True)
        end = self.entity_emb(torch.tensor([ent_ids[-1]]))

        combined = torch.cat([start, rels, end], dim=1)
        return self.path_scorer(combined).squeeze()

    def recommend_with_explanations(self, user_items, kg_adj,
                                    candidates, max_hops=3):
        """返回 [(item_id, score, best_path)]，按分数倒序。"""
        results = []
        for item_id in candidates:
            best_score = float('-inf')
            best_path = None
            for src_item in user_items:
                for path in self.find_paths(kg_adj, src_item, item_id,
                                            max_hops=max_hops):
                    score = self.score_path(path).item()
                    if score > best_score:
                        best_score = score
                        best_path = path
            if best_path is not None:
                results.append((item_id, best_score, best_path))

        results.sort(key=lambda x: x[1], reverse=True)
        return results

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

import torch
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np


class KGRecDataset(Dataset):
    """KG 增强推荐的数据集。"""

    def __init__(self, user_ids, item_ids, ratings):
        self.user_ids = torch.LongTensor(user_ids)
        self.item_ids = torch.LongTensor(item_ids)
        self.ratings = torch.FloatTensor(ratings)

    def __len__(self):
        return len(self.user_ids)

    def __getitem__(self, idx):
        return {
            'user_id': self.user_ids[idx],
            'item_id': self.item_ids[idx],
            'rating': self.ratings[idx],
        }


def train_cke(model, train_loader, val_loader, kg_triples=None,
              num_epochs=10, lr=0.001):
    """训练 CKE 模型，可选地加入 KG 损失。"""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    best_val = float('inf')

    for epoch in range(num_epochs):
        model.train()
        train_loss = 0

        for batch in train_loader:
            optimizer.zero_grad()
            loss = model.compute_loss(
                batch['user_id'], batch['item_id'], batch['rating'],
                *kg_triples if kg_triples else (None, None, None)
            )
            loss.backward()
            optimizer.step()
            train_loss += loss.item()

        model.eval()
        val_loss = 0
        with torch.no_grad():
            for batch in val_loader:
                preds = model(batch['user_id'], batch['item_id'])
                val_loss += F.mse_loss(preds, batch['rating'].float()).item()

        train_loss /= len(train_loader)
        val_loss /= len(val_loader)
        print(f"Epoch {epoch+1}/{num_epochs} | "
              f"Train: {train_loss:.4f} | Val: {val_loss:.4f}")

        if val_loss < best_val:
            best_val = val_loss
            torch.save(model.state_dict(), 'best_kg_model.pt')


if __name__ == "__main__":
    num_users, num_items = 1000, 500
    num_entities, num_relations = 2000, 10

    train_data = KGRecDataset(
        np.random.randint(0, num_users, 8000),
        np.random.randint(0, num_items, 8000),
        np.random.uniform(1, 5, 8000),
    )
    val_data = KGRecDataset(
        np.random.randint(0, num_users, 2000),
        np.random.randint(0, num_items, 2000),
        np.random.uniform(1, 5, 2000),
    )

    model = CKE(num_users, num_items, num_entities, num_relations)
    train_cke(
        model,
        DataLoader(train_data, batch_size=64, shuffle=True),
        DataLoader(val_data, batch_size=64),
        num_epochs=10,
    )