Как я могу распараллелить для l oop для использования в PyTorch?

Question

Я понимаю, что петли for медленны с Python в целом. У меня есть некоторый код, который смешивается с некоторыми тензорами:

            for batch_index, mask_batch in enumerate(mask):
                mask_len = torch.sum(mask_batch).int()

                if mask_len == 0:
                    side_input = torch.zeros((max_inp_len, side_input.shape[1])).to(mask.device)
                else:

                    m_nonzero = mask_batch.nonzero().flatten()
                    first_nonzero = m_nonzero[0]
                    last_nonzero = m_nonzero[-1]

                    if side == 'left':
                        end_index = first_nonzero - 1
                        start_index = 0
                    elif side == 'right':
                        start_index = last_nonzero + 1
                        end_index = inputs[batch_index].size(1)

                    side_input = inputs[batch_index][start_index:end_index]

                    if end_index - start_index < max_inp_len:
                        pad_zeros = torch.zeros(
                            (max_inp_len - side_input.shape[0], side_input.shape[1])).to(mask.device)
                        if side == 'left':
                            side_input = torch.cat((pad_zeros, side_input), 0)
                        elif side == 'right':
                            side_input = torch.cat((side_input, pad_zeros), 0)

                side_inputs.append(side_input)

        return torch.stack(side_inputs)

Мне кажется, что l oop действительно замедляет процесс. Есть ли какой-нибудь способ сделать это без l oop?

Answer 1

Python не имеет истинного параллелизма внутри какого-либо данного процесса. Вам нужно будет создать ProcessPool и сделать внутреннюю часть вашего l oop функцией, принимающей batch_index, mask_batch, а затем отобразить эту функцию на объекте mask в вашем текущем виде для l oop. Дело в том, что я не знаю, будет ли PyTorch хорошо играть с этим.

Примерно так:

def f(batch_index, mask_batch):
    mask_len = torch.sum(mask_batch).int()

    if mask_len == 0:
        side_input = torch.zeros((max_inp_len, side_input.shape[1])).to(mask.device)
    else:
        m_nonzero = mask_batch.nonzero().flatten()
        first_nonzero = m_nonzero[0]
        last_nonzero = m_nonzero[-1]

        if side == 'left':
            end_index = first_nonzero - 1
            start_index = 0
        elif side == 'right':
            start_index = last_nonzero + 1
            end_index = inputs[batch_index].size(1)

            side_input = inputs[batch_index][start_index:end_index]

            if end_index - start_index < max_inp_len:
                pad_zeros = torch.zeros((max_inp_len - side_input.shape[0], side_input.shape[1])).to(mask.device)
                if side == 'left':
                    side_input = torch.cat((pad_zeros, side_input), 0)
                elif side == 'right':
                    side_input = torch.cat((side_input, pad_zeros), 0)
    return side_input

Другие вещи, на которые вы можете посмотреть, - это дальнейшая векторизация кода. Большинство вещей в PyTorch и Numpy можно векторизовать, используя встроенные функции и добавляя в ваши тензоры другое измерение, которое представляет измерение «l oop». Это позволит PyTorch обрабатывать параллелизм за вас.

PyTorch может иметь концепцию устройств, на которые вы можете ставить различные итерации l oop, опять же, для этого вам потребуется создать функцию для этого l oop и, возможно, возьмите устройство, на котором оно работает, в качестве входа.

Наконец, вы можете посмотреть как раз вовремя такие комплименты, как Numba или torch.jit, чтобы выполнить автоматическую векторизацию.

Ничего из этого не будет работать (скорее всего), если mask имеет неизвестную длину. Если он имеет известную длину, я думаю, что векторизация, как бы она ни была сложна, вероятно, является вашим лучшим выбором.

Answer 2

Вы должны создать функцию, содержащую логику c за итерацией al oop, и запустить ее как поток для каждого столбца (см. документы здесь ). Вы также можете использовать библиотеку asyncio для параллелизма, но вы, возможно, получите меньше улучшений.

Хороший пример порождения потока для каждого элемента списка можно прочитать здесь .