torch.cuda.make_graphed_callables

torch.cuda.make_graphed_callables(callables: Module | Callable[[...], object], sample_args: tuple[Tensor, ...], num_warmup_iters: int = 3, allow_unused_input: bool = False, pool: Any | None = None, capture_error_mode: str = 'global', enable_annotations: bool = False) → Module | Callable[[...], object] [исходный код]

torch.cuda.make_graphed_callables(callables:tuple[Module|Callable[[...],object],...], sample_args:tuple[tuple[Tensor,...],...], num_warmup_iters:int=3, allow_unused_input:bool=False, pool:Any|None=None, capture_error_mode:str='global', enable_annotations:bool=False) → tuple[Module|Callable[[...],object],...]

Принимает вызываемые объекты (функции или nn.Module) и возвращает их версии, преобразованные в CUDA-графы.

При прямом проходе каждый преобразованный вызываемый объект выполняет работу исходного вызываемого объекта с CUDA внутри одного узла autograd в виде CUDA-графа.

При прямом проходе преобразованный вызываемый объект также добавляет узел обратного прохода в граф autograd. Во время обратного прохода этот узел выполняет работу вызываемого объекта с CUDA в виде CUDA-графа.

Таким образом, каждый преобразованный вызываемый объект должен быть полноценной заменой исходного вызываемого объекта в цикле обучения с включённым autograd.

Подробные сведения об использовании и ограничениях см. в разделе Захват части сети.

Если передать кортеж из нескольких вызываемых объектов, для их захвата будет использоваться один и тот же пул памяти. О том, в каких случаях это уместно, см. в разделе Управление памятью графа.

Параметры:

• callables (torch.nn.Module или функция Python, или tuple из таких объектов) – Вызываемый объект или объекты для преобразования в CUDA-графы. О том, в каких случаях уместно передавать кортеж вызываемых объектов, см. в разделе Управление памятью графа. Если передаётся кортеж вызываемых объектов, их порядок в кортеже должен совпадать с порядком их выполнения в рабочей нагрузке.
• sample_args (tuple тензоров или tuple кортежей тензоров) – Примеры аргументов для каждого вызываемого объекта. Если передан один вызываемый объект, sample_args должен быть одним кортежем тензоров-аргументов. Если передан кортеж вызываемых объектов, sample_args должен быть кортежем кортежей тензоров-аргументов.
• num_warmup_iters (int) – Количество итераций прогрева. В настоящее время для прогрева DataDistributedParallel требуется 11 итераций. Значение по умолчанию: 3.
• allow_unused_input (bool) – Если значение False, указание входных данных, которые не использовались при вычислении выходных данных (и поэтому их градиент всегда равен нулю), считается ошибкой. Значение по умолчанию — False.
• pool (необязательный) – Токен (возвращаемый функцией graph_pool_handle() или other_Graph_instance.pool()) или MemPool, указывающий, что этот граф может использовать память совместно с указанным пулом. См. раздел Управление памятью графа.
• enable_annotations (bool, необязательный) – Если значение True, при захвате прямого и обратного проходов записываются аннотации ядер из областей torch.cuda.graph_annotations.mark_kernels() внутри вызываемых объектов (ядра обратного прохода помечаются с помощью хуков узлов autograd для этих областей). См. раздел torch.cuda.graph_annotations. Значение по умолчанию: False.

Примечание

Состояние requires_grad каждого тензора в sample_args должно соответствовать состоянию, ожидаемому для соответствующего реального входного значения в цикле обучения.

Предупреждение

Этот API находится на стадии бета-тестирования и может измениться в будущих версиях.

Предупреждение

sample_args каждого вызываемого объекта должен содержать только тензоры. Другие типы не допускаются.

Предупреждение

Возвращаемые вызываемые объекты не поддерживают дифференцирование высших порядков (например, двойное обратное распространение).

Предупреждение

В любом объекте Module, переданном в make_graphed_callables(), обучаемыми могут быть только параметры. Для буферов должно быть установлено requires_grad=False.

Предупреждение

После передачи объекта torch.nn.Module в make_graphed_callables() нельзя добавлять или удалять параметры или буферы этого Module.

Предупреждение

Для объектов torch.nn.Module, передаваемых в make_graphed_callables(), на момент передачи не должны быть зарегистрированы хуки модулей. Однако разрешается регистрировать хуки на модулях после передачи их в make_graphed_callables().

Предупреждение

При вызове преобразованного в граф вызываемого объекта необходимо передавать аргументы в том же порядке и формате, в котором они были указаны в sample_args этого вызываемого объекта.

Предупреждение

Автоматическая смешанная точность поддерживается в make_graphed_callables() только при отключённом кэшировании. Для контекстного менеджера torch.cuda.amp.autocast() должно быть установлено значение cache_enabled=False.