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)

Примеры

>>> 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.])