Воспроизводимая, не зависящая от облака настройка графического процессора и удаленное развертывание кода с помощью Runhouse, пошагового руководства по событию Keras Dreambooth от Hugging Face.

Хакатоны по ИИ — отличный способ поучаствовать в сообществе ИИ и опробовать новые технологии и инструменты для нестандартных вариантов использования; часто эти мероприятия даже сотрудничают с облачными провайдерами, чтобы предоставить участникам бесплатные вычисления! Однако, как бы мы ни стремились погрузиться прямо в код машинного обучения, чтобы начать обучение модели и получение результатов, всегда есть запутанное и трудоемкое препятствие, которое нам нужно преодолеть — настройка нашей удаленной среды разработки.

Буквально в прошлом месяце я участвовал в мероприятии сообщества Hugging Face для тонкой настройки Stable Diffusion (модель преобразования текста в изображение) с использованием Keras Dreambooth, где были предоставлены некоторые бесплатные вычисления LambdaLabs. Несмотря на то, что нам был предоставлен очень полный начальный код для фактического запуска DreamBooth, был довольно длинный список инструкций по настройке удаленных облачных вычислений LambdaLabs, прежде чем мы смогли запустить какой-либо код! Нам нужно было ознакомиться с конкретным пользовательским интерфейсом облачного провайдера (который различается в зависимости от хакатонов), чтобы запускать экземпляры и подключаться к ним по SSH, вручную устанавливать требования и задавать переменные среды, а также передавать данные и код между локальным и удаленным оборудованием.

Runhouse упрощает этот процесс, предоставляя локальный, независимый от облака интерфейс Python для надежного и воспроизводимого запуска и настройки удаленного оборудования и данных. Runhouse помогает преодолеть разрыв между локальной и удаленной разработкой кода в ML, чтобы вам никогда не приходилось вручную подключаться к кластеру по SSH или дважды вводить один и тот же код установки.

Более конкретно, Runhouse позволяет вам

  • Автоматизируйте воспроизводимую настройку оборудования (облачный или локальный экземпляр) с помощью локального кода Python без необходимости запуска экземпляра через облачный пользовательский интерфейс или SSH в экземпляре.
  • Легко синхронизируйте локальные данные и код с удаленными кластерами или хранилищем. Например, получить доступ к локальным тренировочным изображениям в удаленном кластере для предварительной обработки.
  • Запуск локально определенных функций на удаленных кластерах. Кроме того, если у вас есть учетная запись Runhouse (бесплатно!), сохраните и повторно используйте эти функции в локальной и удаленной среде без необходимости их переопределения.

Чтобы увидеть, как Runhouse достигает этого, давайте перейдем к некоторому примеру кода из события DreamBooth Hugging Face Keras и посмотрим, как с помощью всего нескольких дополнительных строк локального кода Runhouse позволяет последовательно и автоматически настраивать и развертывать код в удаленной вычислительной среде. . Полный пример блокнота можно найти здесь.

Установка и настройка Runhouse

Чтобы установить Runhouse: pip install runhouse

Чтобы синхронизировать облачные учетные данные для запуска кластеров AWS, GCP, Azure или LambdaLabs, запустите sky check и следуйте инструкциям. При использовании локального или локального кластера дополнительная настройка не требуется.

Автоматическая настройка оборудования

Runhouse позволяет запускать удаленные облачные экземпляры и надежно настраивать воспроизводимые зависимости и переменные среды с помощью локального кода Python. Больше не нужно вручную запускать кластеры из облачного интерфейса или вводить «стереть данные в экземпляре», чтобы завершить работу вашего экземпляра.

Чтобы развернуть кластер по запросу и создать локальную ссылку на кластер:

import runhouse as rh

gpu = rh.cluster(name="rh-100", instance_type="A100:1", provider="lambda")
gpu.up_if_not()

Теперь, когда кластер настроен и работает, настроить среду так же просто, как gpu.run(list_of_cli_commands) . Runhouse использует gRPC для запуска этого в вашем кластере, поэтому вам не нужно самостоятельно подключаться к машине по SSH.

command = "conda install -y -c conda-forge cudatoolkit=11.2.2 cudnn=8.1.0; \
          mkdir -p $CONDA_PREFIX/etc/conda/activate.d; \
          echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/' > $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh; \
          python3 -m pip install tensorflow"
gpu.run([command])
gpu.restart_grpc_server()  # restart server to use env variables set above

Настройка набора данных

Теперь, чтобы использовать наш локальный набор изображений для точной настройки кластера, мы можем использовать Runhouse для передачи этих изображений на удаленный кластер, используя rh.folder(folder_on_local).to(gpu, folder_on_gpu).

# sync images to the cluster using Runhouse
rh.folder(path=instance_images_root).to(system=gpu, path=instance_images_root)
rh.folder(path=class_images_root).to(system=gpu, path=class_images_root)

Теперь, когда образы находятся в кластере, мы также хотели бы собрать их в кластере. Для этого возьмите нашу локальную функцию assemble_dataset, оберните ее функцией Runhouse и отправьте на наш GPU. Новая функция вызывается так же, как и исходная функция, но магия в том, что она работает на нашем удаленном оборудовании, а не в нашей локальной среде.

assemble_dataset_gpu = rh.function(fn=assemble_dataset).to(system=gpu)

save_data_path = '~/.keras/datasets/train_dataset'
train_dataset_path = assemble_dataset_gpu(new_instance_image_paths, class_image_paths, embedded_text, save_data_path)

Помимо папок, runhouse также позволяет синхронизировать таблицы и хранилище больших двоичных объектов между файловым хранилищем (например, S3) и локальным или удаленным оборудованием.

Обучение

Точно так же мы перемещаем наш обучающий код Dreambooth в функцию train_dreambooth, чтобы мы могли обернуть его как функцию Runhouse для отправки в наш кластер. Как и раньше, мы используем функцию Runhouse, как если бы это была локальная функция, но все обучение происходит на нашем удаленном графическом процессоре.

def train_dreambooth(resolution, max_prompt_length, use_mp, optimizer_params, train_dataset_path, ckpt_path):
  # code to train dreambooth

gpu.run(['cp /usr/lib/cuda/nvvm/libdevice/libdevice.10.bc .'])  # set up libdevice.10.bc to be discoverable by tensorflow
train_dreambooth_gpu = rh.function(fn=train_dreambooth, system=gpu)
ckpt_path_gpu = train_dreambooth_gpu(resolution, max_prompt_length, use_mp, optimizer_params, train_dataset_path, ckpt_path)

Вывод

Ниже мы определяем функцию вывода run_inference и показываем, как мы можем выполнять вывод удаленно, возвращая результаты обратно на локальную сторону для визуализации.

def run_inference(model_path, resolution, max_prompt_length, prompt, num_imgs=3):
   import keras_cv
   from keras_cv.models.stable_diffusion.diffusion_model import DiffusionModel

   # Load in and instantiate model
   sd_dreambooth_model = keras_cv.models.StableDiffusion(img_width=resolution, img_height=resolution, jit_compile=True)
   sd_diffuser_model = DiffusionModel(resolution, resolution, max_prompt_length)
   sd_diffuser_model.load_weights(model_path)
   sd_dreambooth_model._diffusion_model = sd_diffuser_model

   # Run inference
   generated_img = sd_dreambooth_model.text_to_image(prompt, batch_size=num_imgs)
   return generated_img


prompt = f"A photo of two {unique_id} {class_label}"
run_inference_gpu = rh.function(run_inference).to(system=gpu)
inference_images = run_inference_gpu(ckpt_path_gpu, resolution, max_prompt_length, prompt)

Хотя вывод выполняется удаленно, результирующая inference_images является локальной переменной, которую мы можем использовать для локального просмотра результатов! (В качестве альтернативы, если вы предпочитаете оставить результаты в своем кластере, вы можете использовать .remote(), чтобы получить объектную ссылку на изображения в кластере).

Завершение экземпляра

Чтобы завершить экземпляр, как только вы закончите, просто запустите gpu.teardown() в Python или sky down gpu в CLI.

Ресурсы и дополнительная информация

Полную учебную тетрадь с кодом Keras Dreambooth с мероприятия можно найти здесь. Дополнительные примеры использования автонастройки Runhouse для Hugging Face см. в наших интеграциях HF Transformers и Accelerate.