Pytorch параллельный MLP без петли - PullRequest
Новая
       20

Pytorch параллельный MLP без петли

0 голосов
/ 15 января 2020

Вот основной код: 

features = torch.Tensor(feature_num, in_feature)  # feature_num=500,000+
labels = torch.Tensor(feature_num, dtype=torch.long)  # each value range from 0 to label_num-1
weights = torch.Tensor(label_num, in_feature, out_feature)  # in_feature=out_feature=200, label_num=500+

# transform each feature according to its label
mlp_weights = torch.index_select(weights, 0, labels)
out = torch.bmm(features.unsequeeze(dim=1), mlp_weights).unsequeeze(dim=1)

Я просто хочу преобразовать каждый feature, используя спецификацию c weight. Но это вызывает ошибку out of memory , когда выполняется 'index_select' на моей машине с 12 ГБ графического процессора. Я попробовал два решения:

  1. разделить features на несколько кусков, затем сделать то же самое, но это также вызывает ошибку OOM.
  2. l oop в каждой метке и заменить torch.bmm с torch.matmul, но это очень медленно.

Так что любой может дать несколько советов. Спасибо.

...