前言

在阅读yolov5源码时，发现管理进程间数据同步的torch_distributed_zero_first函数用到了上下文管理器，现记录下关于contextmanager和yield结合时的运行逻辑。

@contextmanager
def torch_distributed_zero_first(local_rank: int):
    """
    Decorator to make all processes in distributed training wait for each local_master to do something.
    """
    if local_rank not in [-1, 0]:
        torch.distributed.barrier()
    yield
    if local_rank == 0:
        torch.distributed.barrier()

def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,
                      rank=-1, world_size=1, workers=8, image_weights=False, quad=False, prefix=''):
    # Make sure only the first process in DDP process the dataset first, and the following others can use the cache
    with torch_distributed_zero_first(rank):
        dataset = LoadImagesAndLabels(path, imgsz, batch_size,
                                      augment=augment,  # augment images
                                      hyp=hyp,  # augmentation hyperparameters
                                      rect=rect,  # rectangular training
                                      cache_images=cache,
                                      single_cls=opt.single_cls,
                                      stride=int(stride),
                                      pad=pad,
                                      image_weights=image_weights,
                                      prefix=prefix)

理解

借助一个简单示例对with-as、yield相结合时的运行逻辑进行了理解。

简单示例：

from contextlib import contextmanager

#这段代码的作用是任何对列表的修改只有当所有代码运行完成并且不出现异常的情况下才会生效。
@contextmanager
def list_transaction(orig_list):
    working = list(orig_list)
    yield working
    orig_list[:] = working

items = [1, 2, 3]
with list_transaction(items) as f:
    f.append(6)
    f.append(7)
    raise RuntimeError('oops')

运行逻辑：

1. 在with语句运行后，list_transaction运行到yield部分并抛出函数内的working赋值给f；
2. 在with的作用域内，先向f内新增6、7元素，最后抛出RuntimeError；
3. with作用域完成时，由于提前抛出了RuntimeError，list_transaction函数未能执行orig_list[:] = working。

因为在with作用域内的所有语句执行完后，才会再次进入list_transaction函数执行yield后续的代码，因此with list_transaction(items) as f可以保证其作用域内对items列表的修改只有作用域内代码全部运行完成并且不出现异常的情况下才会生效。

参考

https://python3-cookbook.readthedocs.io/zh_CN/latest/c09/p22_define_context_managers_the_easy_way.html