Блог

Я написал конвейер данных TF, который выглядит примерно так (TF 2.6):

 def parse(img):
    image = tf.image.decode_png(img, channels=3)
    image = tf.reshape(image, IMG_SHAPE)
    image = tf.cast(image, TARGET_DTYPE)
    return image


def decode_batch(serialized_example, is_test=False):
    feature_dict = {
        'image': tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value=''),
    }
    
    if not is_test:
        feature_dict["some_text"] = tf.io.FixedLenFeature(shape=[MAX_LEN], dtype=tf.int64, default_value=[0]*MAX_LEN)
    else:
        feature_dict["image_id"] = tf.io.FixedLenFeature(shape=[], dtype=tf.string, default_value='')

    features = tf.io.parse_example(tf.reshape(serialized_example, [BATCH_SIZE_OVERALL]), features=feature_dict)
    images = tf.map_fn(parse, features['image'], parallel_iterations=4, fn_output_signature=TARGET_DTYPE)

    if is_test:
        image_ids = features["image_id"] 
        return images, image_ids
    else:
        targets = tf.cast(features["some_text"], tf.uint8)
        return images, targets


def get_dataset(filenames, is_test):
    opts = tf.data.Options()
    opts.experimental_deterministic = False
    dataset = tf.data.Dataset.from_tensor_slices(filenames)
    dataset = dataset.with_options(opts)
    dataset = dataset.interleave(lambda x:
        tf.data.TFRecordDataset(x),
        cycle_length=4,
        num_parallel_calls=4,
    )
    dataset = dataset.batch(BATCH_SIZE_OVERALL, num_parallel_calls=4, drop_remainder=True)
    if not is_test:
        dataset = dataset.repeat()
        dataset = dataset.shuffle(BATCH_SIZE_OVERALL*6)
    dataset = dataset.map(lambda y: decode_batch(y, is_test), num_parallel_calls=4)

    dataset = dataset.prefetch(tf.data.AUTOTUNE)
    
    return dataset


train_ds = get_dataset(TRAIN_TFREC_PATHS, False)
 

Как вы можете видеть из кода, я выполнил большинство трюков из руководства TF по правильному построению tf.data конвейера. Проблема у меня следующая: при запуске обучения код использует не все 4 ядра, а только 1 (иногда используется больше ядер, но, похоже, это вызвано train_dist_ds.get_next() вызовом в приведенном ниже коде). Кроме того, графический процессор почти не используется вообще. Профилировщик говорит, что проблема в предварительной обработке, и в tf_data_bottleneck_analysis нем указывается, что проблема в ParallelBatch (хотя однажды он указал на ParallelMap , что кажется правдой, но само по себе это мало о чем говорит — ядра все равно все еще недостаточно используются). Функция обучения с помощью профилировщика выглядит следующим образом:

 def fit_profile(train_ds, val_ds, stop_after_steps):
    tf.profiler.experimental.start('logdir')
    stat_logger.current_step = 0

    train_dist_ds = iter(train_ds)

    while True:
        stat_logger.batch_start_time = time.time()
        stat_logger.current_step  = 1
        print(f'current step: {stat_logger.current_step}')
        with tf.profiler.experimental.Trace('train', step_num=stat_logger.current_step, _r=1):
            image_batch, some_text_batch = train_dist_ds.get_next()
        train_step(image_batch, some_text_batch)
        if stat_logger.current_step == stop_after_steps:
            break
            
    tf.profiler.experimental.stop()
 

Как вы можете видеть, я не прикасаюсь к набору данных, я не включаю его ни в какую стратегию, он train_step есть (в который, конечно, завернут @tf.function ).
Вопросы: есть ли способ каким-то образом отлаживать вычисления внутри графика для tf.data операций? В частности, на уровне вызовов каждой tf.data функции API внутри предварительной обработки-чтобы я мог понять, что именно нужно оптимизировать. В чем может быть причина того, что используется только одно ядро?

То, что я пробовал до сих пор:

• установка всех автоматически настраиваемых параметров на tf.data.AUTOTUNE — без эффекта;
• повторяется только над объектом набора данных-в этом случае используются все ядра, из чего я делаю вывод, что проблема на уровне выполнения графика-параллелизм глобально не отключен;
• выключение профилировщика — никакого эффекта;
• уменьшение количества parallel_iterations входящих map_fn звонков — никакого эффекта;
• множество странных настроек num_parallel_calls — никакого эффекта до такой степени, что кажется, что это действительно не имеет значения.

Я, наконец, нашел причину такого поведения. Это было вызвано использованием XLA с графическим процессором.

Я внезапно обнаружил это и решил отключить XLA, и, о боже, после почти недели исследований графический процессор был полностью использован, а время обучения стало более разумным (до этого они были равны времени обучения процессора!!). Как написано в статье: 1) Поддержка GPU в XLA является экспериментальной; 2) тензоры должны иметь выводимые формы; 3) все операции на графике должны поддерживаться в XLA. Признаками таких проблем являются плохая загрузка ЦП и графического процессора, а также скачкообразные тренировочные шаги, т. е. один шаг занимает 150 секунд, а следующие 8-10 шагов занимают по одной секунде, а затем этот шаблон повторяется. В статье говорится о TF 1.x, но, похоже, до сих пор в этой теме мало что изменилось (опять же, я использую TF 2.6).

Основные приемы:

1. Не используйте XLA с графическим процессором вслепую, это может снизить время обучения вашего графического процессора до уровня процессора (при неправильном использовании).
2. Если вы используете XLA с графическим процессором, убедитесь, что вы соответствуете требованиям, описанным выше.

Я обновлю этот ответ, если мне удастся выполнить эти требования XLA в своих вычислениях и включить XLA с повышением производительности, а не ухудшением.