当前位置: 首页>>代码示例 >>用法及示例精选 >>正文


Python PyTorch MpModelWrapper用法及代码示例


本文简要介绍python语言中 torch_xla.distributed.xla_multiprocessing.MpModelWrapper 的用法。

用法:

class torch_xla.distributed.xla_multiprocessing.MpModelWrapper(model)

当使用fork 方法时,包装一个模型以最小化主机内存使用。

此类应与spawn(…, start_method=’fork’) API 一起使用,以尽量减少主机内存的使用。不是在每个多处理进程上创建模型,从而复制模型的初始主机内存,而是在全局范围内创建一次模型,然后移动到 spawn() 目标函数内的每个设备中。例子:

WRAPPED_MODEL = xmp.MpModelWrapper(MyNetwork())

def _mp_fn(index, ...):
  device = xm.xla_device()
  model = WRAPPED_MODEL.to(device)
  ...

xmp.spawn(_mp_fn, ..., start_method='fork')

这种方法有两个优点。首先,如果仅使用内存页面的一个副本来托管原始模型的权重,其次它通过在此过程中降低系统内存的负载,将包装模型的移动序列化到每个设备中。

相关用法


注:本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torch_xla.distributed.xla_multiprocessing.MpModelWrapper。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。