torch.nn.utils.convert_conv2d_weight_memory_format

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

Преобразует 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. Мы уверены в пользе преобразования NHWC(channels_last) для свёрток в cuDNN, поскольку выполнение свёрток в 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)