Построение распределенных моделей Tensorflow

В этом блоге рассказывается, как легко создавать простые или сложные модели с помощью Tensorflow (≥1.6) Estimator API и Python 3.6+. После создания ряда моделей с использованием Tensorflow Estimator API я понял, что повторно использую много кода, и мои модели обычно приводят к плохому качеству кода. В качестве способа быстрой загрузки моделей, чтобы они быстро заработали, я разработал шаблон, который служит общей основой для большинства моделей. Шаблон позволяет легко проводить локальное или облачное распределенное обучение и пытается включить хорошие стандарты кодирования.

Шаблон во многом вдохновлен https://github.com/MrGemy95/Tensorflow-Project-Template, хотя Estimator API обрабатывает и заменяет многие функции, найденные в этом шаблоне.

Вот ссылка на мой шаблон: https://github.com/maxwellmri/Distributed-Tensorflow-Template, который включает пример MNIST для справки и разбивку шаблона. В репозитории произошел технический сбой.

Требования

Python 3.6+ https://www.python.org/downloads/
git https://git-scm.com/

По желанию

Google Cloud SDK https://cloud.google.com/sdk/docs/ нужен, если вы хотите тренироваться на облачном мл-движке. Примечание: облачные пути к файлам не будут работать в Windows

Примечание: это руководство не охватывает установку tensorflow-gpu, git или gcloud SDK и было протестировано только на Ubuntu 16.04.

Создание новой модели

Клонируйте репозиторий в удобное место из командной строки

git clone https://github.com/maxwellmri/Distributed-Tensorflow-Template.git

2. [Необязательно/Рекомендуется] Создайте виртуальную среду в клонированной папке.

cd Distributed-Tensorflow-Template
virtualenv -p python3.6 example_model_env

2.1 Активировать виртуальную среду

source example_model_env/bin/activate

2.2 Установите тензорный поток, инструменты и пакеты. Обновите requirements.txt, если хотите установить tensorflow-gpu или вам нужны дополнительные пакеты.

pip install -r requirements.txt
pre-commit install

Это включает в себя черный, достойный форматировщик Python, flake8 для линтинга, mypy для дополнительной проверки статического типа и pre-commit, который добавляет хуки перед фиксацией, которые запускают каждый из этих инструментов каждый раз, когда вы совершаете коммит в git ( commit будет работать только в том случае, если все проверки пройдены).

2.3. Обновите удаленный git в свой репозиторий, при необходимости создайте его на Github.

git remote set-url origin https://new.repo-url.here
git add -A
git commit -m "initialised new project from "
git push -u origin master

3. Вкратце, как использовать этот шаблон, например, для реализации модели VGG необходимо сделать следующее:

Обновите сценарий модели, создав класс с именем VGG и унаследовав его от BaseModel.

class VGG16(BaseModel):
    def __init__(self, config: dict) -> None:
        """
        :param config: global configuration
        """
        super().__init__(config)

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

def _create_model(x: tf.Tensor, is_training: bool) -> tf.Tensor:
    """
    :param x: input data
    :param is_training: flag if currently training
    :return: completely constructed model
    """
    pass

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

# TODO: update model predictions
predictions = {
    "classes": tf.argmax(input=logits, axis=1),
    "probabilities": tf.nn.softmax(logits),
}
if mode == tf.estimator.ModeKeys.PREDICT:
    # TODO: update output during serving
    export_outputs = {
        "labels": tf.estimator.export.PredictOutput(
            {"label": predictions["classes"], "id": features["id"]}
        )
    }
    return tf.estimator.EstimatorSpec(
        mode, predictions=predictions, export_outputs=export_outputs
    )

Добавьте любые резюме, которые вы хотите иметь на tensorboard для обучения и оценки.

# TODO: update summaries for tensorboard
tf.summary.scalar("loss", loss)
tf.summary.image("input", tf.reshape(x, [-1, 28, 28, 1]))
# if mode is evaluation
if mode == tf.estimator.ModeKeys.EVAL:
    # TODO: update evaluation metrics
    summaries_dict = {
        "val_accuracy": tf.metrics.accuracy(
            labels, predictions=predictions["classes"]
        )
    }
    return tf.estimator.EstimatorSpec(
        mode=mode, loss=loss, eval_metric_ops=summaries_dict
    )

Измените свой оптимизатор в соответствии с вашей архитектурой, убедитесь, что если вы используете пакетную нормализацию, вы используете зависимости управления (см. пример MNIST).

# TODO: update optimiser
optimizer = tf.train.AdamOptimizer(lr)

Затем обновите сценарий обучения и убедитесь, что он наследуется от класса BaseTrain и использует новую модель и загрузчики данных.

class VGGTrainer(BaseTrain):
    def __init__(
        self,
        config: dict,
        model: VGGModel,
        train: TFRecordDataLoader,
        val: TFRecordDataLoader,
        pred: TFRecordDataLoader,
    ) -> None:
        """
        This function will generally remain unchanged, it is used to 
        train and
        export the model. The only part which may change is the run
        configuration, and possibly which execution to use  
        (training, eval etc)
        :param config: global configuration
        :param model: input function used to initialise model
        :param train: the training dataset
        :param val: the evaluation dataset
        :param pred: the prediction dataset
        """
        super().__init__(config, model, train, val, pred)

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

def _export_model(
        self, estimator: tf.estimator.Estimator, save_location: str
    ) -> None:
        """
        Used to export your model in a format that can be used with
        Tf.Serving
        :param estimator: your estimator function
        """
        # this should match the input shape of your model
        # TODO: update this to your input used in prediction/serving
        x1 = tf.feature_column.numeric_column(
            "input", shape=[self.config["batch_size"], 28, 28, 1]
        )
        # create a list in case you have more than one input
        feature_columns = [x1]
        feature_spec = tf.feature_column.make_parse_example_spec(feature_columns)
        export_input_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
            feature_spec
        )
        # export the saved model
        estimator.export_savedmodel(save_location, export_input_fn)

Обновите загрузчик данных, чтобы корректно загружать входные данные, добавляйте или удаляйте дополнения по мере необходимости. Этот пример для tfrecords, однако можно использовать и другие типы входных данных (https://www.tensorflow.org/guide/datasets). Выполнение загрузки и предварительной обработки данных на ЦП помогает уменьшить узкие места графического процессора.

def _parse_example(
        self, example: tf.Tensor
    ) -> Tuple[Dict[str, tf.Tensor], tf.Tensor]:
        """
        Used to read in a single example from a tf record file and do any augmentations necessary
        :param example: the tfrecord for to read the data from
        :return: a parsed input example and its respective label
        """
        # do parsing on the cpu
        with tf.device("/cpu:0"):
            # define input shapes
            # TODO: update this for your data set
            features = {
                "image": tf.FixedLenFeature(shape=[28, 28, 1], 
                                            dtype=tf.float32),
                "label": tf.FixedLenFeature(shape=[1], 
                                            dtype=tf.int64),
            }
            example = tf.parse_single_example(example, 
                                              features=features)
            if self.mode == "train":
                input_data = self._augment(example["image"])
            else:
                input_data = example["image"]
        return {"input": input_data}, example["label"]

Обновите сценарий задачи, чтобы использовать вашу новую модель и тренажер. Эти сценарии используются для инициализации и обучения вашей модели.

def init() -> None:
    """
    The main function of the project used to initialise all the 
    required functions for training the model
    """
    # get input arguments
    args = get_args()
    # get static config information
    config = process_config()
    # combine both into dictionary
    config = {**config, **args}
    # initialise model
    model = VGGModel(config)
    # create your data generators for each mode
    train_data = TFRecordDataLoader(config, mode="train")
    val_data = TFRecordDataLoader(config, mode="val")
    test_data = TFRecordDataLoader(config, mode="test")
    # initialise the estimator
    trainer = VGGTrainer(config, model, train_data, val_data,  
                         test_data)
    # start training
    trainer.run()

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

def process_config() -> dict:
    """
    Add in any static configuration that is unlikely to change very         
    often
    :return: a dictionary of static configuration data
    """
    config = {"exp_name": "example_model_train"}
    return config

Обучение

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

Файл training_log.md создается или добавляется к нему каждый раз, когда вы запускаете один из сценариев. Вы можете использовать его как рабочий файл для отслеживания деталей эксперимента.

Стандартный вывод вашей модели будет записан в файлы runlogs/, где у каждого соответствующего процесса будет журнал. Он также создает файлы .pid, которые при необходимости можно использовать для завершения процесса. Пример training_log.md показан ниже:

Example Training Job 
Learning Rate: 0.001
Epochs: 100
Batch Size (train/eval): 512/ 512
Hypothesis:
Model will converge quickly
Results:
Model diverged even quicker

Если вам не нравится vim или вы не хотите его использовать, вы можете удалить его из скриптов.

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

##########################################################
# where to write tfevents
OUTPUT_DIR="gs://model-exports"
# experiment settings
TRAIN_BATCH=512
EVAL_BATCH=512
LR=0.001
EPOCHS=100
# create a job name for the this run
prefix="example"
now=$(date +"%Y%m%d_%H_%M_%S")
JOB_NAME="$ENV_NAME"-"$prefix"_"$now"
# locations locally or on the cloud for your files
TRAIN_FILES="data/train.tfrecords"
EVAL_FILES="data/val.tfrecords"
TEST_FILES="data/test.tfrecords"
##########################################################

Обучение на процессоре

Обучение ЦП
Этот скрипт будет обучать модель без использования каких-либо графических процессоров, и вы можете дополнительно указать среду Python для запуска проекта.

использование

Usage: ./train_local_cpu.sh [ENV_NAME]

Обучение на GPU

Обучение графического процессора
Этот скрипт будет обучать модель на конкретном графическом процессоре, и вы можете дополнительно указать среду Python для запуска проекта. Он также проверит, правильно ли вы настроили переменные среды CUDA. Чтобы узнать использование графического процессора или какой GPU_ID использовать, вы можете запустить это в своем терминале.

nvidia-smi

использование

Usage: ./train_local_single.sh <GPU_ID> [ENV_NAME]

Распределенное локальное обучение

Обучение распределенного графического процессора
Этот скрипт позволит вам смоделировать распределенную обучающую среду локально на любом количестве графических процессоров, которое есть на вашей машине. Для этого необходимо разделить GPU на рабочие, главные и серверы параметров. Графические процессоры могут быть отнесены к каждому из этих типов. Вы также можете настроить это прямо в питоне (https://www.tensorflow.org/deploy/distributed). Скрипт содержит пример с использованием 3 графических процессоров:

config="
{
    \"master\": [\"localhost:27182\"],
    \"ps\": [\"localhost:27183\"],
    \"worker\": [
        \"localhost:27184\",
        \"localhost:27185\"
        ]
}, \"environment\": \"cloud\""
...
# ensure parameter server doesn't use any of the GPUs in this case
export CUDA_VISIBLE_DEVICES=""
# Parameter Server can be run on cpu
task="{\"type\": \"ps\", \"index\": 0}"
export TF_CONFIG="{\"cluster\":${config}, \"task\":${task}}"
run ps
# Master should be run on GPU as it runs the evaluation
export CUDA_VISIBLE_DEVICES="1"
task="{\"type\": \"master\", \"index\": 0}"
export TF_CONFIG="{\"cluster\":${config}, \"task\":${task}}"
run master

# Workers (Number of GPUS-1 one used by the master server)
for gpu in 0 1
do
    task="{\"type\": \"worker\", \"index\": $gpu}"
    export TF_CONFIG="{\"cluster\":${config}, \"task\":${task}}"
    export CUDA_VISIBLE_DEVICES="$gpu"
    run "worker${gpu}"
done

В этой установке есть один мастер, один сервер параметров и два рабочих. Мастеру выделяется один графический процессор, а рабочие также имеют по одному графическому процессору. Серверы параметров будут выполняться на ЦП. При определении новых конфигураций необходимо убедиться, что порты, используемые в config, не используются.

использование

Usage: ./train_local_dist.sh [ENV_NAME]

Обучение распределенному облаку

Train Cloud
Этот скрипт требует, чтобы у вас был установлен Google Cloud SDK (https://cloud.google.com/sdk/install) и учетная запись Google Cloud Platform с доступом к ml-engine. Пробные учетные записи GCP поставляются с кредитом, если вы хотите попробовать это. Обучение в облаке стоит денег, но после настройки оно очень простое.

Конфигурация задания

Файл hptuning_config.yaml будет использоваться для указания ресурсов, которые вы запрашиваете для этого задания. Вы можете масштабировать это для своих нужд, оно будет вести себя так же, как локальное распределенное обучение.
Дополнительная информация здесь: https://cloud.google.com/ml-engine/docs/tensorflow/using- gpus
Цены смотрите здесь: https://cloud.google.com/ml-engine/docs/pricing

Требуется, чтобы данные хранились на GCP где-то в ведре, а также нужно указать, в какое ведро экспортировать вашу модель и контрольные точки. Убедитесь, что любые дополнительные пакеты, необходимые вашей модели, определены в setup.py и убедитесь, что вы не указываете пакеты, которые уже являются частью ml-engine (https://cloud.google.com/ml-engine/docs /tensorflow/runtime-список версий)

использование

Usage: ./train_cloud.sh

Просмотр обучения на Tensorboard

Tensorboard поставляется с предварительно установленным tensorflow и является лучшим способом отслеживать ход вашего эксперимента. Вы можете добавлять в свою модель различные показатели, которые полезны для оценки успеха вашей модели. Чтобы запустить tensorboard, убедитесь, что вы используете виртуальную среду или установку python, в которой установлен tensorflow, и запустите его в терминале:

tensorboard --log-dir path_to_your_checkpoints

Путь к вашим контрольным точкам будет определен как один из ваших входных параметров, поэтому, пока вы указываете tensorboard на это место, он автоматически загружает метрики. Это местоположение может быть локальным или путем к корзине GCP (не работает в Windows). Обратите внимание, что тензорная доска, указывающая на корзину GCP, не будет обновляться так часто, как локальный путь.

Заключение

Я надеюсь, что это поможет людям эффективно создавать модели с лучшим качеством кода. Этот шаблон не идеален, и любой вклад в репозиторий приветствуется.