neural factorization machines (SIGIR 2017) unofficial code by Chainer

nfm.py

import numpy as np
import chainer
from chainer import functions as F
from chainer import links as L


class NFMClassifier(chainer.Chain):
    def __init__(self, n_feature, n_dim_emb, n_dim_1, n_dim_2):
        np.random.seed(6162)
        self.n_feature = n_feature
        self.n_dim_emb = n_dim_emb
        self.n_dim_1 = n_dim_1
        self.n_dim_2 = n_dim_2

        self._layers = {
            'EMB': L.EmbedID(self.n_feature, self.n_dim_emb, ignore_label=-1),
            'l_pairwise_1': L.Linear(self.n_dim_emb, self.n_dim_1),
            'l_pairwise_2': L.Linear(self.n_dim_1, self.n_dim_2),
            'l_pairwise_3': L.Linear(self.n_dim_2, 1),
            # sparse fv には先頭に bias term の 1 が入っていることに注意
            'l_linear': L.Linear(self.n_feature + 1, 1),
         }

        super(NFMClassifier, self).__init__(**self._layers)

    # sparse fv には末尾に bias term の 1 が入っていることに注意
    def predict(self, fv, fv_w, sparse_fv):
        # 線形項
        linear = self.l_linear(sparse_fv)

        # 組み合わせを計算していく
        emb = self.EMB(fv)
        emb = emb * fv_w
        pairwise = F.square(F.sum(emb, axis=1)) - F.sum(F.square(emb), axis=1)
        pairwise *= 0.5
        pairwise = F.relu(self.l_pairwise_1(pairwise))
        pairwise = F.relu(self.l_pairwise_2(pairwise))
        pairwise = F.relu(self.l_pairwise_3(pairwise))
        return F.sigmoid(linear + pairwise)

    def __call__(self, fv, fv_w, sparse_fv, y):
        pred = self.predict(fv, fv_w, sparse_fv)

        # RMSE
        loss = F.mean_squared_error(y.reshape(len(y), 1), pred)
        # loss = F.sigmoid_cross_entropy(pred, y.reshape(len(y), 1))
        chainer.report({'loss': loss}, self)

        return loss

train_nfm.py

import argparse
from collections import defaultdict, Counter
import numpy as np
from nfm import NFMClassifier
import chainer
from sklearn.metrics import roc_auc_score

parser = argparse.ArgumentParser(description='Neural Factirozation Machines')
parser.add_argument('--epoch', '-e', default=10, type=int)
parser.add_argument('--dim-1', default=50, type=int)
parser.add_argument('--dim-2', default=50, type=int)
parser.add_argument('--embedding', default=100, type=int)
parser.add_argument('--batchsize', '-b', type=int, default=100)
parser.add_argument('--output-file', default=None)
parser.add_argument('--train-file', default=None)
parser.add_argument('--test-file', default=None)

args = parser.parse_args()
batchsize = args.batchsize
n_epoch = args.epoch
n_dim_1 = args.dim_1
n_dim_2 = args.dim_2
n_emb = args.embedding

# read file
# format : extended libsvm
feature_id_table = defaultdict(lambda: len(feature_id_table))
labels = defaultdict(lambda: list())
fvs = defaultdict(lambda: list())
files = {
    'train': args.train_file,
    'test': args.test_file,
}

for file_type, file in sorted(files.items()):
    with open(file) as f:
        for l in f:
            ary = l.rstrip().split(' ')
            label = int(ary[0])
            fv = { }
            for pair in ary[1:]:
                feature, val = pair.split(':')
                val = float(val)
                feature_id = feature_id_table[feature]
                fv[feature_id] = val

            labels[file_type].append(label)
            fvs[file_type].append(fv)

feature_dim = len(feature_id_table)

model = NFMClassifier(
    n_feature=feature_dim,
    n_dim_emb=n_emb,
    n_dim_1=n_dim_1,
    n_dim_2=n_dim_2
)

# 学習データの数
n_train = len(fvs['train'])
n_test = len(fvs['test'])

# オプティマイザはとりあえず Adam
optimizer = chainer.optimizers.Adam()
optimizer.setup(model)

# trainer に頼る方法がわからなかったので手で minibatch を回す
for epoch in range(1, n_epoch + 1):
    perm = np.random.permutation(n_train)
    sum_loss = 0.0

    minibatch_range = range(0, n_train, batchsize)
    for i in minibatch_range:
        indices = perm[i:i + batchsize]

        # minibatch ごとに特徴量とその重みのベクトルを作る
        minibatch_fvs = [ ]
        minibatch_fv_weights = [ ]
        minibatch_fv_w_mat = [ ]
        minibatch_labels = [ ]
        minibatch_sparse_fvs = np.zeros((len(indices), feature_dim + 1))
        for pos, j in enumerate(indices):
            raw_fv = fvs['train'][j]
            fv = [ ]
            w = [ ]
            # 末尾にバイアス項を追加
            minibatch_sparse_fvs[pos, feature_dim] = 1
            for fv_id, val in sorted(raw_fv.items()):
                fv.append(fv_id)
                w.append(val)
                minibatch_sparse_fvs[pos, fv_id] = val

            minibatch_fvs.append(fv)
            minibatch_fv_weights.append(w)
            minibatch_labels.append(labels['train'][j])

        # 最も多く発火していた fv に合わせて -1 で pad する
        maximum_activated_fv_size = max([len(fv) for fv in minibatch_fvs])

        # -1 で padding
        for pos in range(len(indices)):
            old_fv = minibatch_fvs[pos]
            pad_length = maximum_activated_fv_size - len(old_fv)

            minibatch_fvs[pos] = np.pad(
                old_fv, (0, pad_length),  mode='constant', constant_values=-1
            )

            # 重みについては埋め込み次元だけ伸ばしたベクトルを発火した分だけ必要
            # w が一次元のベクトルから行列になる
            fv_w = np.zeros((maximum_activated_fv_size, n_emb))
            for _pos, w in enumerate(minibatch_fv_weights[pos]):
                fv_w[_pos, ] = w

            minibatch_fv_w_mat.append(fv_w)

        minibatch_fvs = chainer.Variable(np.asarray(minibatch_fvs, dtype=np.int32))
        minibatch_fv_weights = chainer.Variable(np.asarray(minibatch_fv_w_mat, dtype=np.float32))
        minibatch_labels = chainer.Variable(np.asarray(minibatch_labels, dtype=np.float32))
        minibatch_sparse_fvs = chainer.Variable(np.asarray(minibatch_sparse_fvs, dtype=np.float32))

        model.cleargrads()
        loss = model(minibatch_fvs, minibatch_fv_weights, minibatch_sparse_fvs, minibatch_labels)
        loss.backward()
        optimizer.update()
        sum_loss += loss.data

    print(sum_loss)

    # 予測
    predicts = [ ]
    test_indices = np.array(range(n_test))
    for i in range(0, n_test, batchsize):
        indices = test_indices[i:i + batchsize]
        # minibatch ごとに特徴量とその重みのベクトルを作る
        minibatch_fvs = [ ]
        minibatch_fv_weights = [ ]
        minibatch_fv_w_mat = [ ]
        minibatch_sparse_fvs = np.zeros((len(indices), feature_dim + 1))
        for pos, j in enumerate(indices):
            raw_fv = fvs['test'][j]
            fv = [ ]
            w = [ ]
            # 末尾にバイアス項を追加
            minibatch_sparse_fvs[pos, feature_dim] = 1
            for fv_id, val in sorted(raw_fv.items()):
                fv.append(fv_id)
                w.append(val)
                minibatch_sparse_fvs[pos, fv_id] = val

            minibatch_fvs.append(fv)
            minibatch_fv_weights.append(w)

        # 本来なら全特徴量次元で -1 するわけですがそれでは無駄が多い
        # なのである minibatch において最も多く発火していた fv に合わせて -1 で pad する
        maximum_activated_fv_size = max([len(fv) for fv in minibatch_fvs])

        # -1 で padding
        for pos in range(len(indices)):
            old_fv = minibatch_fvs[pos]
            pad_length = maximum_activated_fv_size - len(old_fv)

            minibatch_fvs[pos] = np.pad(
                old_fv, (0, pad_length),  mode='constant', constant_values=-1
            )

            # 重みについては埋め込み次元だけ伸ばしたベクトルを発火した分だけ必要
            # w が一次元のベクトルから行列になる
            fv_w = np.zeros((maximum_activated_fv_size, n_emb))
            for _pos, w in enumerate(minibatch_fv_weights[pos]):
                fv_w[_pos, ] = w

            minibatch_fv_w_mat.append(fv_w)

        minibatch_fvs = chainer.Variable(np.asarray(minibatch_fvs, dtype=np.int32))
        minibatch_fv_weights = chainer.Variable(np.asarray(minibatch_fv_w_mat, dtype=np.float32))
        minibatch_sparse_fvs = chainer.Variable(np.asarray(minibatch_sparse_fvs, dtype=np.float32))

        raw_predicts = model.predict(minibatch_fvs, minibatch_fv_weights, minibatch_sparse_fvs)
        predicts += [v[0] for v in raw_predicts.data]

    # calc auc
    print(roc_auc_score(labels['test'], predicts))