- from contextlib import contextmanager
-
- @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]:
- dist.barrier(device_ids=[local_rank])
- yield
- if local_rank == 0:
- dist.barrier(device_ids=[0])
-
-
-
- with torch_distributed_zero_first(rank):
- dataset = LoadImagesAndLabels(path)
通过上下文管理器以便使用with语句。使用 contexlib
模块中的 @contextmanager
装饰器。上下文管理器通过with语句来自动获取和释放资源,确保资源的正确管理和关闭,避免资源泄露和错误。
上下文管理器中有__enter__()和__exit__()俩个特殊方法。__enter__()方法在with语句块执行前调用,负责获取资源并返回资源对象。__exit__()在with语句执行完调用,负责资源的释放。
创建上下文管理器有两种方法:
(1)通过函数装饰器创建,使用 contexlib
模块中的 @contextmanager
装饰器。
- from contextlib import contextmanager
-
- @contextmanager
- def my_context():
- try:
- print('Entering the context') # 获取资源,例如打开文件、建立数据库连接等
- yield # yield语句之前的代码相当于__enter__()方法
- # yield语句之后的代码相当于__exit__()方法# 使用函数装饰器创建上下文管理器
-
- except Exception as e:
- print(f'An exception of type {type(e)} occurred with value {e}')
-
- finally:
- print('Exiting the context') # 释放资源,例如关闭文件、关闭数据库连接等
-
-
- with my_context():
- print('Inside the context') # 在这里使用资源进行操作
运行结果:
装饰器修饰my_context()函数使它成为一个上下文管理器。yield语句用于指示资源的获取及释放,允许在with语句中使用资源。
(2)通过类创建
- class MyContent:
- def __enter__(self):
- print('Entering the content') # 获取资源,例如打开文件等
- return self
-
- def __exit__(self, exc_type, exc_val, exc_tb):
- print('Exiting the content') # 释放资源,例如关闭文件
- if exc_type is not None:
- print(f'An exception of type {exc_type} occurred with value {exc_val}')
- return False # 返回False表示不处理异常,异常会继续传播
-
- # 使用自定义的上下文管理器
- with MyContent() as my_content:
- print('Inside the content') # 在这里使用资源对象进行操作print('Outside the context')
Mycontent()是一个自定义的上下文管理器。__enter__()方法返回self自身作为资源对象。__exit__()在with语句执行完之后被自动调用,释放资源。
上下文管理器的使用:具备上下文管理器之后,可以直接使用with语句管理资源的获取和释放。
with 上下文管理器 as 资源对象:
对资源对象的操作
with语句开始执行时,会先调用__enter__()方法中的内容用于获取资源,这一部分可以进行资源对象的操作。with执行完毕之后,无论是否发生异常,都会被调用用于释放资源。
- class FileHandler:
- def __init__(self, filename, mode):
- self.filename = filename
- self.mode = mode
-
- # 打开文件并返回文件对象,可以直接使用file对象进行文件的读取操作
- def __enter__(self):
- self.file = open(self.filename, self.mode)
- return self.file
-
- # with执行完毕,调用__exit__()关闭文件
- def __exit__(self):
- self.file.close()
-
- # 上下文管理器
- with FileHandler('file.txt', 'r') as file: #
- content = file.read()
- print(content) # __exit__()中关闭文件