torch.nn.utils.prune.custom_from_mask

torch.nn.utils.prune.custom_from_mask(module, name, mask) [source]

Обрезает тензор, соответствующий параметру, называемому name в module, применяя предварительно вычисленную маску в mask. Модифицирует модуль на месте (и также возвращает изменённый модуль) следующим образом:

1. добавляет именованный буфер, называемый name+'_mask', соответствующий двоичной маске, применённой к параметру name методом обрезки.
2. заменяет параметр name его обрезанной версией, в то время как исходный (необрезанный) параметр сохраняется в новом параметре, названном name+'_orig'.

Параметры

• module (nn.Module) – модуль, содержащий тензор для обрезки
• name (str) – имя параметра внутри module, к которому будет применена обрезка.
• mask (Tensor) – двоичная маска для применения к параметру.

Возвращает

изменённую (т.е. обрезанную) версию входного модуля

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

модуль (nn.Module)

Примеры

>>> from torch.nn.utils import prune
>>> m = prune.custom_from_mask(
...     nn.Linear(5, 3), name='bias', mask=torch.tensor([0, 1, 0])
... )
>>> print(m.bias_mask)
tensor([0., 1., 0.])