torch.utils.module_tracker

Создано: 4 мая 2024 г. | Последнее обновление: 11 июня 2025 г.

Эту утилиту можно использовать для отслеживания текущего положения внутри иерархии torch.nn.Module. Её можно использовать в других средствах отслеживания, чтобы легко связывать измеренные величины с понятными пользователю именами. В настоящее время она используется, в частности, в FlopCounterMode.

class torch.utils.module_tracker.ModuleTracker [исходный код]

ModuleTracker — это менеджер контекста, который отслеживает иерархию nn.Module во время выполнения, чтобы другие системы могли запрашивать, какой Module выполняется в данный момент (или выполняется ли его обратный проход).

Чтобы получить набор всех Modules, выполняемых в данный момент, в виде их fqn (полностью квалифицированного имени, также используемого в качестве ключа в state_dict), можно обратиться к атрибуту parents этого менеджера контекста. Чтобы узнать, выполняется ли в данный момент обратный проход, можно обратиться к атрибуту is_bw.

Обратите внимание, что parents никогда не бывает пустым и всегда содержит ключ “Global”. Флаг is_bw остаётся True после прямого прохода, пока не будет выполнен другой Module. Если вам нужна более точная информация, отправьте запрос на добавление этой возможности. Добавление отображения fqn в экземпляр модуля возможно, но пока не реализовано; если вам это необходимо, отправьте соответствующий запрос.

Пример использования

mod = torch.nn.Linear(2, 2)

with ModuleTracker() as tracker:
    # Access anything during the forward pass
    def my_linear(m1, m2, bias):
        print(f"Current modules: {tracker.parents}")
        return torch.mm(m1, m2.t()) + bias

    torch.nn.functional.linear = my_linear

    mod(torch.rand(2, 2))