torch.nn.utils.convert_conv3d_weight_memory_format

torch.nn.utils.convert_conv3d_weight_memory_format(module, memory_format) [исходный код]

Преобразует memory_format из nn.Conv3d.weight в memory_format. Преобразование рекурсивно применяется к вложенным nn.Module, включая module. Обратите внимание, что оно меняет только memory_format, но не семантику каждого измерения. Эта функция используется для того, чтобы задействовать ядра NHWC, которые обеспечивают значительное ускорение при обработке данных fp16 на устройствах CUDA с вычислительной способностью >= 7.0.

Примечание

Вызов model.to(memory_format=torch.channels_last_3d) — более агрессивное действие, чем использование вспомогательной функции convert_conv3d_weight_memory_format. model.to затрагивает любой слой с весами размерности 4, что не обязательно дает преимущества при преобразовании в указанный memory_format. Мы уверены в целесообразности преобразования NDHWC(channels_last_3d) для свертки в cuDNN, поскольку выполнение свертки в NDHWC выгодно даже в случаях, когда требуется переставлять входные тензоры.

Поэтому здесь мы используем стратегию преобразования только весов свертки в channels_last_3d. Это гарантирует, что: 1. Будут использоваться быстрые ядра свертки, преимущества которых могут перевесить накладные расходы на перестановку (если входные данные имеют другой формат). 2. Не будут выполняться ненужные перестановки в слоях, которым не дает преимуществ преобразование memory_format.

Оптимальный вариант — когда слои между слоями свертки совместимы с форматом channels_last. При прохождении через первый слой свертки входной тензор будет переставлен в формат channels_last и останется в этом формате памяти. Поэтому последующим слоям свертки не потребуется переставлять входной тензор.

Если между слоями свертки находится слой, несовместимый с форматом channels_last, потребуется вернуть входной тензор в непрерывный формат для этого слоя. В оставшихся слоях входной тензор будет проходить в непрерывном формате и будет переставлен в формат channels_last при прохождении через следующий слой свертки. Нет смысла распространять эту перестановку на более ранний слой, поскольку большинство слоев довольно невосприимчивы к memory_format.

Это утверждение может измениться, когда PyTorch будет поддерживать слияние перестановок, поскольку, возможно, найдется более подходящее место для слияния перестановки, чем непосредственно перед сверткой.

Параметры:

• module (nn.Module) – nn.Conv3d & nn.ConvTranspose3d или контейнер nn.Module
• memory_format (memory_format) – заданный пользователем memory_format, например torch.channels_last или torch.contiguous_format

Возвращает:

Исходный модуль с обновленным nn.Conv3d

Тип возвращаемого значения:

_M

Пример

>>> input = torch.randint(
...     1, 10, (2, 8, 4, 4, 4), dtype=torch.float16, device="cuda"
... )
>>> model = nn.Sequential(
>>>     nn.Conv3d(8, 4, 3)).cuda().half()
>>> # This is identical to:
>>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d)
>>> model = nn.utils.convert_conv3d_weight_memory_format(
...     model, torch.channels_last_3d
... )
>>> out = model(input)