如何使用 sklearn 中的 train_test_split 确保用户和项目同时出现在训练和测试数据集中？

Question

I have a data set including user ID , item ID , and rating as below:

user ID     item ID    rating
 1233        1011       4
 1220        0999       3
 2011        0702       1
 ...

When I split them into train and test sets:

from sklearn import cross_validation

train, test = cross_validation.train_test_split(df, test_size = 0.2)

Whether the users in test set have already appeared in the train set, and so have items? If not, how can I do that? I can not find the answer in document . Could you please tell me?

Answer 1

If you want to ensure that your training and test partitions don't contain the same pairings of user and item then you could replace each unique (user, item) combination with an integer label, then pass these labels to LabelKFold . To assign integer labels to each unique pairing you could use this trick :

import numpy as np
import pandas as pd
from sklearn.cross_validation import LabelKFold

df = pd.DataFrame({'users':[0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2],
                   'items':[0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1],
                   'ratings':[2, 4, 3, 1, 4, 3, 0, 0, 0, 1, 0, 1]})

users_items = df[['users', 'items']].values
d = np.dtype((np.void, users_items.dtype.itemsize * users_items.shape[1]))
_, uidx = np.unique(np.ascontiguousarray(users_items).view(d), return_inverse=True)

for train, test in LabelKFold(uidx):

    # train your classifier using df.loc[train, ['users', 'items']] and
    # df.loc[train, 'ratings']...

    # cross-validate on df.loc[test, ['users', 'items']] and
    # df.loc[test, 'ratings']...

---

I'm still having a hard time understanding your question. If you want to guarantee that your training and test sets do contain examples of the same user then you could use StratifiedKFold :

for train, test in StratifiedKFold(df['users']):
    # ...

Answer 2

def train_test_split(self, ratings, train_rate=0.8):
        """
        Split ratings into Training set and Test set

        """
        grps = ratings.groupby('user_id').groups
        test_df_index = list()
        train_df_index = list()

        test_iid = list()
        train_iid = list()

        for key in grps:
            count = 0
            local_index = list()
            grp = np.array(list(grps[key]))

            n_test = int(len(grp) * (1 - train_rate))
            for i, index in enumerate(grp):
                if count >= n_test:
                    break
                if ratings.iloc[index]['movie_id'] in test_iid:
                    continue
                test_iid.append(ratings.iloc[index]['movie_id'])
                test_df_index.append(index)
                local_index.append(i)
                count += 1

            grp = np.delete(grp, local_index)

            if count < n_test:
                local_index = list()
                for i, index in enumerate(grp):
                    if count >= n_test:
                        break
                    test_iid.append(ratings.iloc[index]['movie_id'])
                    test_df_index.append(index)
                    local_index.append(i)
                    count += 1

                grp = np.delete(grp, local_index)

            train_df_index.append(grp)

        test_df_index = np.hstack(np.array(test_df_index))
        train_df_index = np.hstack(np.array(train_df_index))

        np.random.shuffle(test_df_index)
        np.random.shuffle(train_df_index)

        return ratings.iloc[train_df_index], ratings.iloc[test_df_index]

You can use this method to split, I've already done efforts to make sure that the training set and test set have the same user id and movie id.