かえるのプログラミングブログ

プログラミングでつまずいたところとその解決策などを書いていきます。

単語 ID 列を長さでソートしてミニバッチ内で padding する。

こんばんは、kaerururu (@kaeru_nantoka) です。

今回は、kaggle meetup #6 での tks さんの発表にもありました、「 (batch 内で) batch 毎に padding する」の実装にプラスして ID列を長さでソートしたものを batch 内で padding できるようにした実装を公開しようと思います。

誤り、もっと良い実装等ございましたらご指摘いただけると幸いです。m( )m

何が嬉しいの?

ID 列の長さでソートしたものを dataloader に渡し、batch 内で padding すると 初めの方は、系列長 0 や 1, 2 などが渡ってくるので padding した後の長さが 1 や 2 になり無駄がなくなるので処理速度が爆速になります。

昨今の kaggle の NLP コンペ (NLPに限った話ではないですが..) では、学習は local でもやっても良いが推論は kaggle kernel 内で済ませる必要があるタイプが多く、いかに個別の model にかかる時間を節約し、規定の推論時間内に model を詰め込みアンサンブルするかが勝負の分かれ目になります。

デメリット

model の training 時には学習にかかる時間を大きく減らす効果が期待できます。一方で、以下の論文  

https://www.anlp.jp/proceedings/annual_meeting/2017/pdf_dir/A7-1.pdf

でも指摘されているように、長さの似た文章は似た情報を持っていることも考えられるため同じミニバッチに詰め込んで学習させることで model が偏った知識を学習してしまうことで精度が少し落ちてしまう恐れもあります。

実装

「 batch 毎に 0 padding して系列長を揃える 」までは、以下の kernel の実装を参考にしました。

https://www.kaggle.com/kunwar31/pytorch-pad-sequences-per-batch/notebook

あとは系列の長さで sort した x_train (sorted_x_train) とそれに対応する順番に並べ替えた y_train (sorted_y_train) を dataset に渡すだけで OK です。

sort した x_train と y_train を入手する関数を以下のように用意しました。

[[長さ], [x_train], [y_train]] のように配列を作り、1番目の配列(長さ)でソートしたのち、単体の配列に詰め直して返してます。

from operator import itemgetter

def get_sorted_list(lengths, x_train, y_train=None):
    tmp_list = []
    if y_train is not None:
        for l, m, n in zip(lengths, x_train, y_train):
            tmp_list.append([l, m, n])  
        tmp_list.sort(key=itemgetter(0))
        sorted_x_train = []
        sorted_y_train = []
        for i in tmp_list:
            sorted_x_train.append(i[1]) 
            sorted_y_train.append(i[2]) 
        return sorted_x_train, sorted_y_train
    else:
        for l, m in zip(lengths, x_train):
            tmp_list.append([l, m])  
        tmp_list.sort(key=itemgetter(0))
        sorted_x_train = []
        for i in tmp_list:
            sorted_x_train.append(i[1]) 
        return sorted_x_train

class TextDataset(data.Dataset):
    '''
    Simple Dataset
    '''
    def __init__(self,X,y=None):
        self.X = X
        self.y = y

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

    def __getitem__(self, idx):
        if self.y is not None:
            return [self.X[idx],self.y[idx]]
        return self.X[idx]


class MyCollator(object):
    '''
    Yields a batch from a list of Items
    Args:
    test : Set True when using with test data loader. Defaults to False
    percentile : Trim sequences by this percentile
    '''
    def __init__(self,test=False,percentile=100):
        self.test = test
        self.percentile = percentile
    def __call__(self, batch):
        if not self.test:
            data = [item[0] for item in batch]
            target = [item[1] for item in batch]
        else:
            data = batch
        lens = [len(x) for x in data]
        max_len = np.percentile(lens,self.percentile)
        data = sequence.pad_sequences(data,maxlen=int(max_len))
        data = torch.tensor(data,dtype=torch.long).cuda()
        if not self.test:
            target = torch.tensor(target,dtype=torch.float32).cuda()
            return [data,target]
        return [data]
collate = MyCollator(percentile=100)

lengths = torch.from_numpy(np.array([len(x) for x in x_train]))

sorted_x_train, sorted_y_train = get_sorted_list(lengths, x_train, y_train_final)

train_dataset = TextDataset(sorted_x_train, sorted_y_train)

train_loader  = torch.utils.data.DataLoader(train_dataset, batch_size=512, shuffle=False, collate_fn=collate)

まとめ

Quora や jigsaw の solution を見ていて実装したいと思っていたのですが、日本語での解説記事など探しても見当たらなかったので公開しようと思った次第です。

(batch padding sort kaggle でググって公開カーネルにたどり着いた)

余力があれば、

1 ) 精度面でどれだけ悪化するのか

2 ) どれだけ早くなるのか

を比較実験した記事も書きたいなと思います。

ありがとうございました。