torch.backends.mkl 的源代码

import torch


def is_available():
    r"""返回 PyTorch 是否使用 MKL 支持构建。"""
    return torch._C.has_mkl


VERBOSE_OFF = 0
VERBOSE_ON = 1


class verbose:
    """
    按需的 oneMKL 详细功能。

    为了更容易调试性能问题，oneMKL 可以转储详细消息，包含执行信息，如执行内核时的持续时间。详细功能可以通过名为 `MKL_VERBOSE` 的环境变量调用。然而，这种方法会在所有步骤中转储消息。这些是大量的详细消息。此外，为了调查性能问题，通常只需要一次迭代的详细消息就足够了。此按需详细功能使得可以控制详细消息转储的范围。在以下示例中，详细消息将仅在第二次推理时转储。

    .. highlight:: python
    .. code-block:: python

        import torch
        model(data)
        with torch.backends.mkl.verbose(torch.backends.mkl.VERBOSE_ON):
            model(data)

    参数:
        level: 详细级别
            - ``VERBOSE_OFF``: 禁用详细
            - ``VERBOSE_ON``:  启用详细
    """

    def __init__(self, enable):
        self.enable = enable

    def __enter__(self):
        if self.enable == VERBOSE_OFF:
            return
        st = torch._C._verbose.mkl_set_verbose(self.enable)
        assert (
            st
        ), "Failed to set MKL into verbose mode. Please consider to disable this verbose scope."
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        torch._C._verbose.mkl_set_verbose(VERBOSE_OFF)
        return False