Wskazówki dotyczące wydajności

Ten dokument zawiera wskazówki dotyczące wydajności dla zestawów danych TensorFlow (TFDS). Należy zauważyć, że TFDS udostępnia zestawy danych jako obiekty tf.data.Dataset , więc porady zawarte w przewodniku tf.data nadal mają zastosowanie.

Porównawcze zbiory danych

Użyj tfds.benchmark(ds) do testu porównawczego dowolnego obiektu tf.data.Dataset .

Upewnij się, że wskazano parametr batch_size= , aby znormalizować wyniki (np. 100 iter/s -> 3200 ex/sec). Działa to z każdą iteracją (np tfds.benchmark(tfds.as_numpy(ds)) ).

ds = tfds.load('mnist', split='train').batch(32).prefetch()
# Display some benchmark statistics
tfds.benchmark(ds, batch_size=32)
# Second iteration is much faster, due to auto-caching
tfds.benchmark(ds, batch_size=32)

Małe zbiory danych (mniej niż 1 GB)

Wszystkie zestawy danych TFDS przechowują dane na dysku w formacie TFRecord . W przypadku małych zbiorów danych (np. MNIST, CIFAR-10/-100), odczyt z .tfrecord może dodać znaczne obciążenie.

Ponieważ te zestawy danych mieszczą się w pamięci, można znacznie poprawić wydajność przez buforowanie lub wstępne ładowanie zestawu danych. Zauważ, że TFDS automatycznie buforuje małe zestawy danych (poniższa sekcja zawiera szczegóły).

Buforowanie zbioru danych

Oto przykład potoku danych, który jawnie buforuje zestaw danych po znormalizowaniu obrazów.

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label


ds, ds_info = tfds.load(
    'mnist',
    split='train',
    as_supervised=True,  # returns `(img, label)` instead of dict(image=, ...)
    with_info=True,
)
# Applying normalization before `ds.cache()` to re-use it.
# Note: Random transformations (e.g. images augmentations) should be applied
# after both `ds.cache()` (to avoid caching randomness) and `ds.batch()` (for
# vectorization [1]).
ds = ds.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.cache()
# For true randomness, we set the shuffle buffer to the full dataset size.
ds = ds.shuffle(ds_info.splits['train'].num_examples)
# Batch after shuffling to get unique batches at each epoch.
ds = ds.batch(128)
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)

Podczas iteracji tego zestawu danych druga iteracja będzie znacznie szybsza niż pierwsza dzięki buforowaniu.

Automatyczne buforowanie

Domyślnie automatyczne pamięci podręczne TFDS (z ds.cache() ) spełniają następujące ograniczenia:

  • Całkowity rozmiar zbioru danych (wszystkie podziały) jest zdefiniowany i < 250 MiB
  • shuffle_files jest wyłączone lub odczytywany jest tylko jeden fragment

Można zrezygnować z automatycznego buforowania, podając try_autocaching=False do tfds.ReadConfig w tfds.load . Zapoznaj się z dokumentacją katalogu zestawów danych, aby sprawdzić, czy określony zestaw danych będzie korzystał z automatycznego buforowania.

Ładowanie pełnych danych jako jeden Tensor

Jeśli Twój zestaw danych mieści się w pamięci, możesz również załadować pełny zestaw danych jako pojedynczą tablicę Tensor lub NumPy. Można to zrobić, ustawiając batch_size=-1 , aby wsadowo wszystkie przykłady w jednym tf.Tensor . Następnie użyj tfds.as_numpy do konwersji z tf.Tensor na np.array .

(img_train, label_train), (img_test, label_test) = tfds.as_numpy(tfds.load(
    'mnist',
    split=['train', 'test'],
    batch_size=-1,
    as_supervised=True,
))

Duże zbiory danych

Duże zestawy danych są dzielone na fragmenty (dzielone na wiele plików) i zazwyczaj nie mieszczą się w pamięci, więc nie powinny być buforowane.

Tasowanie i szkolenie

Podczas treningu ważne jest, aby dobrze przetasować dane — źle potasowane dane mogą skutkować niższą dokładnością treningu.

Oprócz używania ds.shuffle do tasowania rekordów należy również ustawić shuffle_files=True , aby uzyskać dobre zachowanie tasowania dla większych zestawów danych, które są podzielone na wiele plików. W przeciwnym razie epoki będą czytać odłamki w tej samej kolejności, a więc dane nie będą naprawdę losowe.

ds = tfds.load('imagenet2012', split='train', shuffle_files=True)

Dodatkowo, gdy shuffle_files=True , TFDS wyłącza options.deterministic , co może nieznacznie zwiększyć wydajność. Aby uzyskać deterministyczne tasowanie, można zrezygnować z tej funkcji za pomocą tfds.ReadConfig : albo przez ustawienie read_config.shuffle_seed lub nadpisanie read_config.options.deterministic .

Automatyczne dzielenie danych między pracownikami (TF)

Podczas uczenia wielu procesów roboczych można użyć argumentu input_context tfds.ReadConfig , aby każdy pracownik odczytał podzbiór danych.

input_context = tf.distribute.InputContext(
    input_pipeline_id=1,  # Worker id
    num_input_pipelines=4,  # Total number of workers
)
read_config = tfds.ReadConfig(
    input_context=input_context,
)
ds = tfds.load('dataset', split='train', read_config=read_config)

Jest to uzupełnienie interfejsu API subsplit. Najpierw stosowany jest subplit API: train[:50%] jest konwertowany na listę plików do odczytu. Następnie do tych plików stosowana jest ds.shard() . Na przykład przy użyciu train[:50%] z num_input_pipelines=2 , każdy z 2 pracowników odczyta 1/4 danych.

Gdy shuffle_files=True , pliki są tasowane w obrębie jednego pracownika, ale nie między pracownikami. Każdy pracownik będzie czytać ten sam podzbiór plików między epokami.

Automatyczne dzielenie danych między pracownikami (Jax)

Dzięki Jaxowi możesz użyć interfejsu API tfds.split_for_jax_process lub tfds.even_splits do dystrybucji danych między pracownikami. Zobacz dzielony przewodnik po interfejsach API .

split = tfds.split_for_jax_process('train', drop_remainder=True)
ds = tfds.load('my_dataset', split=split)

tfds.split_for_jax_process to prosty alias dla:

# The current `process_index` loads only `1 / process_count` of the data.
splits = tfds.even_splits('train', n=jax.process_count(), drop_remainder=True)
split = splits[jax.process_index()]

Szybsze dekodowanie obrazu

Domyślnie TFDS automatycznie dekoduje obrazy. Istnieją jednak przypadki, w których bardziej wydajne może być pominięcie dekodowania obrazu za pomocą tfds.decode.SkipDecoding i ręczne zastosowanie tf.io.decode_image op:

Kod dla obu przykładów jest dostępny w przewodniku dekodowania .

Pomiń nieużywane funkcje

Jeśli używasz tylko podzbioru funkcji, możesz całkowicie pominąć niektóre funkcje. Jeśli Twój zbiór danych zawiera wiele nieużywanych funkcji, ich nieodkodowanie może znacznie poprawić wydajność. Zobacz https://www.tensorflow.org/datasets/decode#only_decode_a_sub-set_of_the_features