咨询：基于~300GB数据集训练Convnet（Keras/TensorFlow）的最佳实践

Hey there! Handling large-scale datasets like yours (300GB subset of a 15TB corpus) for semantic segmentation with Keras/TensorFlow requires some optimized workflows to avoid bottlenecks and memory issues. Here are the best practices I’ve refined from production projects and research work:

The biggest bottleneck with large datasets is almost always data loading—here’s how to fix it:

• Ditch Keras generators for tf.data API: Keras generators are slower and less efficient for big data. Instead, use tf.data.Dataset.list_files to load file paths, then write a custom parsing function to load images and segmentation masks on-the-fly. Use dataset.map(your_parse_fn, num_parallel_calls=tf.data.AUTOTUNE) to parallelize loading, and finish with dataset.prefetch(tf.data.AUTOTUNE) to overlap data loading and model training.
• Use TFRecord format for faster I/O: Convert your raw images/masks into TFRecord files (split into 1-5GB shards). TFRecords are optimized for TensorFlow, reducing disk seek time and making data loading faster. You can write a simple script to batch-process your dataset into these shards, then load them with tf.data.TFRecordDataset.
• Lazy loading only: Never load the entire dataset into memory. All processing (decoding, resizing, augmentation) should happen per-sample as part of the data pipeline. Use tf.io.read_file to load raw bytes, then tf.image.decode_jpeg/png to decode images—this keeps memory usage low.

• Mixed precision training: Enable mixed precision with tf.keras.mixed_precision.set_global_policy('mixed_float16'). This uses 16-bit floats for most computations while keeping 32-bit weights, cutting memory usage by ~50% and speeding up training (especially on GPUs with Tensor Cores).
• Gradient accumulation: If your GPU can’t handle a large batch size, split your batch into smaller mini-batches, accumulate gradients over several steps, then update weights once. You can implement this with a custom callback or use built-in utilities in newer TensorFlow versions.
• Gradient checkpointing: For large convnets (like U-Net variants with many layers), enable gradient checkpointing. This trades a small amount of training time for reduced memory usage by recomputing some forward-pass values during backprop instead of storing them. Use tf.recompute_grad on your model layers or set model.compile(run_eagerly=False) to leverage this.

If you have access to multiple GPUs or TPUs, don’t waste their potential:

• Multi-GPU with MirroredStrategy: Wrap your model definition and compilation in a MirroredStrategy scope. This automatically splits data across GPUs, runs computations in parallel, and aggregates gradients. Example:

  strategy = tf.distribute.MirroredStrategy()
  with strategy.scope():
      model = build_your_semantic_segmentation_model()
      model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
  

• TPU training (if available): TPUs are built for massive-scale workloads. Use tf.distribute.TPUStrategy to leverage their parallel processing power—they handle large datasets and compute-heavy convnets far more efficiently than GPUs in many cases.

• Integrate augmentation into the tf.data pipeline: Use TensorFlow’s native image augmentation functions (like tf.image.random_flip_left_right, tf.image.random_brightness, tf.image.random_crop) directly in your parsing function. This ensures augmentation runs on GPU (or parallel CPU threads) instead of being a bottleneck.
• Random patch cropping: Instead of using full-size images, crop random fixed-size patches (e.g., 256x256 or 512x512) during training. This reduces memory usage, increases data diversity, and helps your model learn finer-grained features—critical for semantic segmentation.

• Weighted sampling: If some classes are underrepresented, use tf.data.Dataset.sample_from_datasets to assign higher sampling weights to rare classes. Alternatively, pass a class_weight dictionary to model.fit() to adjust loss contributions.
• Focus on hard examples: You can also use loss functions that prioritize hard-to-segment regions (like focal loss) to make your model pay more attention to minority classes.

内容的提问来源于stack exchange，提问作者Arun Aniyan