Декораторы контекста — это конкретная реализация менеджеров контекста Python. Я проиллюстрирую, как их использовать, на примере (pytorch) отладки графического процессора, где я нахожу их весьма полезными. Это не сработает во всех ситуациях, но в целом я хочу поделиться тем, как написание ваших собственных контекстных менеджеров может стать отличным инструментом для вашего набора навыков в Python, и чем они отличаются от декораторов функций.
тлдр; вот суть github с кодом
Проблема с утечкой памяти
Существует множество способов отладки утечек памяти. Здесь я покажу один подход, который я считаю полезным при определении проблемной строки в коде. Подход помогает найти это аккуратным способом.
Построчная ручная отладка
Типичным подходом грубой силы будет проверка построчно с помощью отладчика. Если вы уже знаете, где искать, то это может быть хорошим подходом.
- Google для фрагмента о том, как посчитать все тензоры в pytorch
- Установить точку останова в коде
- Проверьте количество тензоров с помощью фрагмента
- Сделайте один шаг с отладчиком
- Повторно запустите tensor-counter-snippet и проверьте, увеличилось ли количество тензоров.
- Повторить
Это немного хлопотно. Вы также можете поместить тензорный подсчет в функцию и скопировать и вставить его по всему коду, чтобы распечатать любые увеличения. Это работает, но что было бы более оптимальным, так это один фрагмент, чтобы сделать все это с гибкой длиной кода, с небольшой необходимостью изменять существующие строки кода.
Декораторы Python
Одним из шагов в правильном направлении может быть декоратор, который оборачивает любую произвольную часть кода и проверяет, не созданы ли какие-либо дополнительные тензоры. Это можно сделать с помощью декораторов. Часто упоминаемый пример — синхронизация функции, но в данном случае мы будем считать тензоры до и после выполнения. Шаблон выглядит следующим образом:
def memleak_wrapper(func):
def wrap(*args, **kwargs):
print("num tensors start is ...")
out = func(*args, **kwargs)
print("num tensors end is ...")
return out
return wrap
@memleak_wrapper
def function_to_debug(x):
print(f"put line(s) of code here. Input is {x}")
out = x + 10
return out
out = function_to_debug(x=1000)
print(f"out is {out}")
вывод будет:
num tensors start is ... put line(s) of code here. Input is 1000 num tensors end is ... outis 1010
Чтобы заставить его работать, вам нужно поместить строки кода, которые вы хотите проверить, в функцию (function_to_debug). Это не оптимально, потому что требует кучу ручного труда. Кроме того, если блок кода генерирует больше переменных, чем одна 1, вам нужно будет найти дополнительные решения для их использования ниже по течению.
Декораторы контекста:
Вместо декоратора функций можно использовать контекстный менеджер. Наиболее широко используемый пример контекстных менеджеров — создание экземпляра контекста с помощью оператора with. Вы видели их много раз раньше:
with open(“file”) as f: …
Пользователи Python могут легко создавать контекстные менеджеры самостоятельно, используя contextlib, используя библиотеку contextlib Python. В этом случае мы собираемся использовать ContextDecorator, который делает то же, что мы пытались делать с декоратором выше, но более просто. Нам нужно импортировать его и написать класс, который наследуется от него, вот так:
from contextlib import ContextDecorator
class check_memory_leak_context(ContextDecorator):
def __enter__(self):
print('Starting')
return self
def __exit__(self, *exc):
print('Finishing')
return False
У него есть 2 метода: __enter__() и __exit__() , которые вызываются при входе или выходе из контекста.
Аргумент *exc в __exit__ обозначает любые входящие исключения, к которым вы можете получить доступ оттуда.
return selfв __enter__, если это полезно, если вы хотите использовать экземпляр f, как при записи в него в open():
with open(“file”) as f: f.write(‘HELLO!’)
Теперь для нашей реализации утечки памяти нам это не нужно:
Использование ContextDecorator для обнаружения утечки памяти
Мы будем только get_n_tensors() подсчитывать количество тензоров в начале и конце контекста:
class check_memory_leak_context(ContextDecorator):
def __enter__(self):
self.start = get_n_tensors()
return self
def __exit__(self, *exc):
self.end = get_n_tensors()
increase = self.end — self.start
if increase > 0:
print(f”num tensors increased with"\
f"{self.end — self.start} !”)
else:
print(”no added tensors”)
return False
Если есть какое-то увеличение, мы выводим это на консоль.
get_n_tensors() использует сборщик мусора (gc) и настроен для pytorch, но может быть легко изменен для любой библиотеки. Вот как это выглядит:
import gc
def get_n_tensors():
tensors= []
for obj in gc.get_objects():
try:
if (torch.is_tensor(obj) or
(hasattr(obj, ‘data’) and
torch.is_tensor(obj.data))):
tensors.append(obj)
except:
pass
return len(tensors)
Теперь все, что нам нужно сделать, это использовать этот контекст для любой строки (или блоков) кода:
x = arbitrary_operation(x)
...
with check_memory_leak_context():
y = x[0].permute(1, 2, 0).cpu().detach().numpy()
x = some_harmless_operation()
...
x = another_arbitrary_operation(x)
Приведенный выше код является произвольным, но он будет напечатан независимо от того, будут ли строки, украшенные контекстом, создавать новый тензор или нет.
Я обычно держу этот фрагмент под рукой или в моем репозитории в отдельном файле во время разработки. Вот GitHub суть с кодом
Надеюсь, что этот небольшой пост дал некоторое представление о том, что такое менеджеры контекста, как использовать декораторы контекста и как применять их для отладки pytorch.
следуй за мной в твиттере, чтобы быть в курсе любых будущих сообщений
