Блог

Главная - Вопросы по программированию - Создание системы очередей pytorch выполняется в Google Colab. Как я могу надежно очистить память CUDA без перезапуска среды выполнения?

Создание системы очередей pytorch выполняется в Google Colab. Как я могу надежно очистить память CUDA без перезапуска среды выполнения?

#pytorch #google-colaboratory #generative-adversarial-network

Вопрос:

Я работаю с версией клипа VQGAN , которая должна позволить мне создавать очередь поколений изображений pytorch в Google Colab. Между каждым запуском я получал RuntimeError сообщение об ошибке из-за нехватки памяти CUDA. Я мог бы перезапустить время выполнения ноутбука, и он запустился бы и работал нормально.

Чтобы создать очередь запусков, я добавил эту строку, чтобы очистить кэш CUDA для повторного запуска: torch.cuda.empty_cache()

С этой строкой, добавленной в итерацию по поколениям изображений, я решил, что она очистит кэш всякий раз, когда столкнется с этим RuntimeError . Однако, похоже, это не работает, и перезапуск среды выполнения, по-видимому, является единственным способом начать запуск.

Вот цикл, который запускает обучение pytorch:

 def train_and_update(i, last_image=False, retryTime=0):
    try:
        new_filepath = self.train(i, fileargs["Base_Dir"], last_image)
        pbar.update()
        return new_filepath

    except RuntimeError as e:
        print("RuntimeError: "   str(e))
        if retryTime>0:
            print(f"Retrtying in {retryTime}.nYou may need to lower your size settings or change models.")
        torch.cuda.empty_cache()
        time.sleep(retryTime)
        newRetryTime = retryTime   2
        train_and_update(i, last_image=last_image, retryTime=newRetryTime)

Репозиторий кода записной книжки находится здесь: https://github.com/Hoversquid/VQGAN_CLIP_Z_Quantize

Настоящая записная книжка здесь: https://colab.research.google.com/drive/1djy-WUawvcetuAqWfoEOsPRz_FrpKsCd

Должна ли очистка кэша CUDA устранить эту проблему, или я ошибаюсь в том, что torch.cuda.empty_cache() следует делать?

Метки: