検索アルゴリズムを作りたい

Tweet

pytorchでRankNet

import torch
import torch.nn as nn
import torch.optim as optim

# モデル定義
class RankNet(nn.Module):
    def __init__(self, input_dim):
        super(RankNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_dim, 10),
            nn.ReLU(),
            nn.Linear(10, 1)
        )

    def forward(self, x1, x2):
        score1 = self.fc(x1)
        score2 = self.fc(x2)
        return score1, score2

# 損失関数定義
def ranknet_loss(score1, score2, label):
    # label = 1 if doc1 is more relevant than doc2, else 0
    diff = score1 - score2
    return torch.log(1 + torch.exp(-label * diff)).mean()

# データとモデルの設定
input_dim = 5
model = RankNet(input_dim)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# トイデータ
# x1, x2 は文書の特徴ベクトル, labelはx1がx2よりも関連性が高い場合は1、そうでない場合は0
x1 = torch.tensor([[2.0, 3.0, 1.0, 5.0, 4.0], [2.0, 4.0, 1.0, 5.0, 3.0]])
x2 = torch.tensor([[1.0, 2.0, 2.0, 4.0, 5.0], [1.0, 1.0, 3.0, 4.0, 5.0]])
labels = torch.tensor([1, 0], dtype=torch.float32)

# トレーニング
for epoch in range(100):
    optimizer.zero_grad()

    score1, score2 = model(x1, x2)
    loss = ranknet_loss(score1, score2, labels)

    loss.backward()
    optimizer.step()

    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1}, Loss: {loss.item()}")

print("Training finished!")

このサンプルコードでは、2つの文書ペアの特征ベクトルと、それらの関連性のラベルを基に、RankNetモデルを訓練しています。実際の応用では、もちろんもっと多くのデータと、適切な特徴ベクトルを使用する必要があります。