hfai.utils¶

`which_numa`	根据 gpu 编号或者 ib 编号得到对应的 numa
`num_gpus`	返回可用的 GPU 个数
`profile_memory`	分析模型的显存占用情况
`bind_numa`	绑定当前进程到指定的 NUMA 节点
`get_current_numa`	返回当前进程所在的 NUMA 编号
`timeout`	给函数设置超时限制

hfai.utils.which_numa(i_gpu=None, i_ib=None)[source]¶

根据 gpu 编号或者 ib 编号得到对应的 numa

Parameters

i_gpu (int, optional) – gpu 编号
i_ib (int, optional) – ib 编号

Returns

对应的 numa 编号

Return type

int

Examples

>>> from hfai.utils import which_numa
>>> which_numa(i_gpu=0)
0

hfai.utils.num_gpus()[source]¶

返回可用的 GPU 个数

相比于 torch.cuda.device_count()，hfai.utils.num_gpus 不用初始化 cuda，对使用 fork 启动训练任务的用户更加友好

Returns: 可用的 GPU 个数
Return type: int

Examples

>>> import os, hfai
>>> hfai.utils.num_gpus()
8
>>> os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
>>> hfai.utils.num_gpus()
4

hfai.utils.profile_memory(model, input=(), input_kwargs={}, include_children=True, sort_by='name', show_shapes=False, show_peakmem=False, show_forward_time=False, forward_funcs=['forward'])[source]¶

分析模型的显存占用情况

打印出来的结果包含以下几个字段：

parameter size: 参数总量
activation size: forward 的过程中通过 save_for_backward 保存的 tensor 大小（不包含参数）
#calls: 被调用的次数
input shape: 输入的 tensor 形状
output shape: 输出的 tensor 形状
peak mem: 峰值显存；forward 过程中的峰值显存减去 forward 之前的已占用显存
forward time: 模块时间：每个模块　forward　前后的时间差，多次调用则累加

Note

不同算子可能会重复保存一部分的中间层变量，所以总的 activation size 会比实际的显存使用量要大。

Note

show_peakmem = True 和 include_children = False 互斥

Note

show_forward_time = True 和 include_children = False 互斥

Note

仅支持 PyTorch >= 1.10

Parameters

model (torch.nn.Module) – 需要被分析的模型
input (tuple) – 模型的输入，通过 model(*input, **input_kwargs) 调用
input_kwargs (dict) – 模型的关键字参数，通过 model(*input, **input_kwargs) 调用
include_children (bool) – 每个模块的显存占用计算是否包含其子模块（类型为 nn.Module）的显存占用；默认是 True
sort_by (str) – name, activation, parameter, peakmem 或者 forward_time，输出的时候根据哪个字段进行排序；默认是 name
show_shapes (bool) – 是否打印输入、输出的形状；默认是 False
show_peakmem (bool) – 是否打印峰值显存，模型必须在 GPU 上；默认是 False
show_forward_time (bool) – 是否打印模块时间

Examples

>>> import torch, hfai
>>> from torchvision import models
>>> model = models.alexnet().cuda()
>>> x = torch.randn(64, 3, 224, 224, device="cuda")
>>> hfai.utils.profile_memory(model, input=(x,))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
HFAI Memory Profiler (include_children = True, sort_by = name)
====================  =================  ================  =================  ========
module name           type                 parameter size    activation size    #calls
====================  =================  ================  =================  ========
AlexNet               AlexNet                 233.081 MiB        345.312 MiB         1
AlexNet.features      Sequential                9.421 MiB        336.000 MiB         1
AlexNet.features.0    Conv2d                    0.089 MiB         36.750 MiB         1
......
AlexNet.classifier.5  ReLU                      0.000 MiB          1.000 MiB         1
AlexNet.classifier.6  Linear                   15.629 MiB          1.000 MiB         1
====================  =================  ================  =================  ========
total unique activations: 225.906 MiB
======================================================================================

hfai.utils.bind_numa(node)[source]¶

绑定当前进程到指定的 NUMA 节点

Parameters: node (int) – NUMA 节点编号

Examples

>>> import hfai
>>> hfai.utils.bind_numa(0)
>>> hfai.utils.get_current_numa()
0
>>> hfai.utils.bind_numa(1)
>>> hfai.utils.get_current_numa()
1

hfai.utils.get_current_numa()[source]¶

返回当前进程所在的 NUMA 编号

Examples

>>> import hfai
>>> hfai.utils.bind_numa(0)
>>> hfai.utils.get_current_numa()
0
>>> hfai.utils.bind_numa(1)
>>> hfai.utils.get_current_numa()
1

hfai.utils.timeout(seconds=0)[source]¶

给函数设置超时限制

Parameters: seconds (int) – 超时的秒数

Examples:

import time, hfai
from hfai.utils import TimeoutError

@hfai.utils.timeout(10)
def print_hello():
    while True:
        print("hello")
        time.sleep(1)

try:
    print_hello()
except TimeoutError as e:
    print(e)