Как использовать глобальный тензор в параллельном режиме данных pytorch?

#mpi #pytorch

#mpi #pytorch

Вопрос:

В этом примере я бы хотел, чтобы z_proto мог быть глобальным для разных графических процессоров. Однако в режиме параллельного анализа данных он также разделяется на разные графические процессоры. Как решить такую проблему? Спасибо.

 class SequencePrototypeTokenClassification(nn.Module):
    def __init__(self,seq_model, label_num):
        super(SequencePrototypeTokenClassification, self).__init__()
        self.seq_model = seq_model
        self.label_num = label_num

    def forward(self, input_ids, token_type_ids, attention_mask, labels, z_proto, n_query, target_inds):
        z, _ = self.seq_model(input_ids, token_type_ids, attention_mask, output_all_encoded_layers=False)
        z_dim = z.size(-1)
        zq = z.squeeze().view(-1, z_dim)
        dists = euclidean_dist(zq, z_proto)
        log_p_y = F.log_softmax(-dists, dim=1).view(-1, self.label_num)
        loss_val = -log_p_y.gather(1, self.target_inds).squeeze().view(-1).mean()
        _, y_hat = log_p_y.max(1)

        return loss_val, y_hat
  

Ответ №1:

Основываясь на вашем приведенном выше коде, z_proto кажется, является одним из аргументов прямой функции, а не частью модели. Следовательно, простое сохранение его в tensor на главном графическом процессоре позволит ему иметь одинаковое значение для всех графических процессоров.

Редактировать

Судя по документации, кажется, что DataParallel все входные данные для функции прямой передачи распределяются по графическим процессорам. Метод, с помощью которого вы можете обойти это, — сохранить его как переменную класса внутри самого объекта модели. Вы можете обновить значение перед вызовом функции пересылки, если это не статическая переменная.

 class SequencePrototypeTokenClassification(nn.Module):
    def __init__(self,seq_model, label_num):
        ...
        self.z_proto = None
        ...
        ...


#Training loop
    ...
    model.z_proto = value
    model.forward()
    ...


  

Комментарии:

1. Я поставил z_proto на основной графический процессор. Но replicas = self.replicate(self.module, self.device_ids[:len(inputs)]) в DataParallel это разделило бы z_proto на 4 графических процессора.

2. Привет, я попробовал этот метод. proto_model.module.z_proto = z_proto Но он получил RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generated/../generic/THCTensorMathPointwise.cu:272

3. Это странно. Согласно docs , pytorch выполняет разделение только во время прямого вызова и объединяет его обратно перед следующей строкой. Вы уверены, что разместили z_proto обновления за пределами forward() функции? А вы инициализировались z_proto в __init__ функции?

4. Да, я сделал почти то же самое, что и ты. Я инициализировал z_proto с None помощью forward и обновил его раньше,,, но используя только GPU 0. Я думал, z_proto будет воспроизведен с module , но этого не произошло.

5. Можете ли вы поделиться своим полным кодом вместе с циклом обучения?

Ответ №2:

Оказывается, DataParallel будет реплицировать только nn.Parameter из nn.Module . Итак, я случайным образом инициализировал nn.Parameter имя z_proto в модуле и скопировал значение тензора z_proto в параметр. Затем параметр реплицируется на 4 графических процессора.