torch.nn.utils.memory_format.convert_conv2d_weight_memory_format

torch.nn.utils.memory_format.convert_conv2d_weight_memory_format(module, memory_format) [source]

Преобразует memory_format nn.Conv2d.weight в memory_format.

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

Примечание

Вызов model.to(memory_format=torch.channels_last) выполняет более агрессивное преобразование, чем служебная функция convert_conv2d_weight_memory_format. model.to затронет любой слой с весами размерности 4, что не обязательно принесёт пользу от преобразования в указанный memory_format. В одном случае мы уверены в пользе: преобразование свёртки в cuDNN в NHWC(channels_last), поскольку выполнение свёртки в NHWC полезно даже в тех случаях, когда необходимо применять перестановку к входным тензорам.

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

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

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

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

Параметры:

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

Возвращает:

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

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

_M

Пример

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