Время прочтения: 3 мин.

Если Вам приходится постоянно обучать модели, будь то Machine Learning, или задачи в области Computer Vision, искать и классифицировать какие-либо объекты, то Вы знаете, что ожидание результата и многочисленные итерации занимают безумно много времени. Хотите сократить время обучения и повысить эффективность работы?

Как правило при решении задач машинного обучения используется только одно вычислительное устройство, например, GPU. Но многие модели обучаются на видеокарте столько же времени, как и на процессоре. Я расскажу, как распараллелить процессы для реализации обучения моделей одновременно на GPU и CPU.

Создадим класс для запуска обучения на процессоре. Он будет выглядеть следующим образом:

class CPUManager:
    def __init__(self):
        self.env = os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
    def trainer(self, train_test_data, model):
        res_dict = {}
        """
        Тут ваш код
        """
        return res_dict

В методе “__init__” указываем для ‘CUDA_VISIBLE_DEVICES’ значение ‘-1’, чтобы графический процессор не был использован, так как если несколько процессов будут пытаться использовать одно и то же устройство, то это может привести к ошибке OutOfMemoryError.

В методе “trainer” прописываем свою реализацию обучения модели или код для ее применения на тестовых данных.

Создадим класс для запуска обучения на графическом процессоре. Он будет выглядеть следующим образом:

class GPUManager:
    def __init__(self):
        self.env = os.environ['CUDA_VISIBLE_DEVICES'] = "0"
    def trainer(self, train_test_data, model):
        res_dict = {}
        """
        Тут ваш код
        """
        return res_dict

В методе “__init__” указываем для ‘CUDA_VISIBLE_DEVICES’ значение ‘0’, чтобы выбрать доступный графический процессор.

В методе “trainer” так же прописываем реализацию обучения модели или, как сказано выше, код для ее применения.

Перейдем к основному классу, который запускает наши модели и отвечает за их распараллеливание.

Для сбора итога выполнения нашего обучения или использования моделей использую класс Manager библиотеки Multiprocessing. Это реализация паттерна «заместитель», предназначенная для возможности использовать общий объект несколькими запущенными процессами. Используем его для сбора результатов выполнения кода применимых моделей так как это удобнее, чем отлавливать результаты каждого процесса по отдельности.

Используем “run”, как основной метод класса, для создания двух процессов. Первый будет принимать функцию “CPU”, указываем её как target и передаем саму модель и путь до общего датасета. Точно так же делаем со вторым методом, только передаем ему функцию “GPU”. Эти функции создают объекты своих классов, отвечающие за их реализации, после запускается метод “trained”, который также принимает датасет и саму модель. Сохраняем в наш прокси-объект результат выполнения функции и ожидаем окончание работы всех процессов.

class ParallelTeachModel:
    def __init__(self):
        self.manager = multiprocessing.Manager()
        self.return_dict = self.manager.dict()

    def CPU(self, data, model):
        """
        Этот метод выполняется с использованием заданного пути до датасета упомянутой модели.
        Аргументы:
        data (str): путь до датасета
        model (object): CPU модель, которую хотим использовать
        """
        obj_cpu_manager = CPUManager()
        output_res = obj_cpu_manager.trainer(data, model)
        self.return_dict['CPU_output'] = output_res

    def GPU(self, data, model):
        """
        Этот метод выполняется с использованием заданного пути до датасета упомянутой модели.
        Аргументы:
        data (str): путь до датасета
        model (object): GPU модель, которую хотим использовать
        """
        obj_gpu_manager = GPUManager()
        output_res = obj_gpu_manager.trainer(data, model)
        self.return_dict['GPU_output'] = output_res

    def run(self, dataset, cpu_model, gpu_model):
        """
        Это основной код, который выполняет обработку.
        Аргументы:
        data (str): путь до датасета
        cpu_model (object): CPU модель, которую хотим использовать
        gpu_model (object): GPU модель, которую хотим использовать
        Результат:
        dict: результат обработки
        """
        # запуск CPU модели
        process1 = multiprocessing.Process(target=self.CPU, args=(dataset, cpu_model,))
        process1.start()

        # запуск GPU модели
        process2 = multiprocessing.Process(target=self.GPU, args=(dataset, gpu_model,))
        process2.start()

        # Объединение процессов
        process1.join()
        process2.join()

        return self.return_dict

Запускаем! Создаем экземпляр класса “Parallel_Teach_Model”, указываем путь до датасета, прописываем модели для CPU и GPU, после вызываем функцию “run” того же класса. Функция возвращает выходные данные моделей в виде словаря, который мы можем использовать дальше.

if __name__ == '__main__':
    obj = ParallelTeachModel()
    dataset_path = "path"
    cpu_model = object()
    gpu_model = object()
    preds = obj.run(dataset_path, cpu_model, gpu_model)
print(preds)

В данной статье мы разобрали способ оптимизации обучения или использования моделей, который необходим для сокращения времени обучения и повышения эффективности работы.