torch.nn.utils.memory_format.convert_conv3d_weight_memory_format

torch.nn.utils.memory_format.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. На любой слой с весами размерности 4D повлияет model.to, хотя преобразование в указанный 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)