CIFAR-100データセットの一部クラスのデータだけを使いたい

CIFAR-10またはCIFAR-100データセットの一部のクラスのデータセットのみを抽出して、使用したいのですが方法がわからないです。例えば、５クラス分のみを抽出し、５クラス分類問題として学習と評価をしたいです。
全データを使う場合は、以下のコードで大丈夫なのですが、そこからどのように変更すれば良いのでしょうか。

python
1import cPickle
2import numpy as np
3import os
4
5def unpickle(file):
6    fo = open(file, 'rb')
7    dict = cPickle.load(fo)
8    fo.close()
9    return dict
10
11def conv_data2image(data):
12    return np.rollaxis(data.reshape((3,32,32)),0,3)
13
14def get_cifar10(folder):
15    tr_data = np.empty((0,32*32*3))
16    tr_labels = np.empty(1)
17    '''
18    32x32x3
19    '''
20    for i in range(1,6):
21        fname = os.path.join(folder, "%s%d" % ("data_batch_", i))
22        data_dict = unpickle(fname)
23        if i == 1:
24            tr_data = data_dict['data']
25            tr_labels = data_dict['labels']
26        else:
27            tr_data = np.vstack((tr_data, data_dict['data']))
28            tr_labels = np.hstack((tr_labels, data_dict['labels']))
29
30    data_dict = unpickle(os.path.join(folder, 'test_batch'))
31    te_data = data_dict['data']
32    te_labels = np.array(data_dict['labels'])
33
34    bm = unpickle(os.path.join(folder, 'batches.meta'))
35    label_names = bm['label_names']
36    return tr_data, tr_labels, te_data, te_labels, label_names
37
38def get_cifar100(folder):
39    train_fname = os.path.join(folder,'train')
40    test_fname  = os.path.join(folder,'test')
41    data_dict = unpickle(train_fname)
42    train_data = data_dict['data']
43    train_fine_labels = data_dict['fine_labels']
44    train_coarse_labels = data_dict['coarse_labels']
45
46    data_dict = unpickle(test_fname)
47    test_data = data_dict['data']
48    test_fine_labels = data_dict['fine_labels']
49    test_coarse_labels = data_dict['coarse_labels']
50
51    bm = unpickle(os.path.join(folder, 'meta'))
52    clabel_names = bm['coarse_label_names']
53    flabel_names = bm['fine_label_names']
54
55    return train_data, np.array(train_coarse_labels), np.array(train_fine_labels), test_data, np.array(test_coarse_labels), np.array(test_fine_labels), clabel_names, flabel_names
56
57if __name__ == '__main__':
58    datapath = "./data/cifar-10-batches-py"
59    datapath2 = "./data/cifar-100-python"
60
61    tr_data10, tr_labels10, te_data10, te_labels10, label_names10 = get_cifar10(datapath)
62    tr_data100, tr_clabels100, tr_flabels100, te_data100, te_clabels100, te_flabels100, clabel_names100, flabel_names100 = get_cifar100(datapath2)

行動規範の内容に同意します

回答1件

ベストアンサー

とりあえず、CIFAR-10 の方法のみ記述します。（試してませんが基本的に CIFAR-100 でも同じかと思います）

Python
1# 0,1,2,3,4 のみを対象とする
2target = [0,1,2,3,4]
3
4tr_mask = np.isin(tr_labels10 , target)
5custom_tr_data10 = tr_data10[tr_mask, :]
6custom_tr_labels10 = tr_labels10[tr_mask]
7
8te_mask = np.isin(te_labels10, target)
9custom_te_data10 = te_data10[te_mask, :]
10custom_te_labels10 = te_labels10[te_mask]