如何减少张量流数据集输入管道主机设备cpu时间目前约40％？

如何解决如何减少张量流数据集输入管道主机设备cpu时间目前约40％？

我正在尝试复制resnet18论文。在磁盘上的完整Image Net数据集上运行此代码之前，我要对TFDS中公开可用的imagenette/320px数据集（包含10个类的imagenet小得多的子集，已经采用.tfrecord格式）进行一些评估._

注意：可以在此处进行培训和跟踪的完整笔记本：resnet18_baseline.ipynb只需切换到GPU运行时并运行所有单元。它已经在第二批中使用张量板配置进行设置。 （您也可以使用TPU，但是某些keras.layers.experimental.preprocessing层尚不支持TPU操作，并且您必须启用软设备放置。请使用GPU。）

输入操作

1. 从输入数据集中读取图像。这些图像通常具有不同的尺寸，并且我们需要一些裁剪功能，因为输入张量不能不具有不同的尺寸用于批处理。因此，在训练中，我使用随机作物，而在测试/验证数据集中使用中心作物。

random_crop_layer = keras.layers.experimental.preprocessing.RandomCrop(224,224)
center_crop_layer = keras.layers.experimental.preprocessing.CenterCrop(224,224)

@tf.function(experimental_relax_shapes=True) # avoid retracing
def train_crop_fn(x,y):
  return random_crop_layer(x),y

@tf.function(experimental_relax_shapes=True)
def eval_crop_fn(x,y):
  return center_crop_layer(x),y

1. 对输入数据进行一些简单的预处理/扩充。其中包括重新缩放为0-1，以及基于imagenet上rgb颜色的均值和stdev进行缩放。另外，随机

rescaling_layer = keras.layers.experimental.preprocessing.Rescaling(1./255)

train_preproc = keras.Sequential([
  rescaling_layer
])

# from https://github.com/tensorflow/models/blob/master/official/vision/image_classification/preprocessing.py
# Calculated from the ImageNet training set
MEAN_RGB = (0.485,0.456,0.406)
STDDEV_RGB = (0.229,0.224,0.225)

@tf.function
def z_score_scale(x):
  return (x - MEAN_RGB) / STDDEV_RGB

@tf.function
def train_preproc_fn(x,y):
  return z_score_scale(train_preproc(x)),y

@tf.function
def eval_preproc_fn(x,y):
  return z_score_scale(eval_preproc(x)),y

输入管道

def get_input_pipeline(input_ds,bs,crop_fn,augmentation_fn):
  ret_ds = (
      input_ds
      .batch(1) # pre-crop are different dimensions and can't be batched
      .map(crop_fn,num_parallel_calls=tf.data.experimental.AUTOTUNE)
      .unbatch()
      .batch(bs)
      .map(augmentation_fn,# augmentations can be batched though.
          num_parallel_calls=tf.data.experimental.AUTOTUNE)
  )

  return ret_ds

# dataset loading
def load_imagenette():
  train_ds,ds_info = tfds.load('imagenette/320px',split='train',as_supervised=True,with_info=True)
  valid_ds = tfds.load('imagenette/320px',split='validation',as_supervised=True)

  return train_ds,valid_ds,ds_info.features['label'].num_classes

# pipeline construction
train_ds,test_ds,num_classes = load_imagenette()

# datasets used for training (notice that I use prefetch here)
train_samples = get_input_pipeline(train_ds,BS,train_crop_fn,train_preproc_fn).prefetch(tf.data.experimental.AUTOTUNE)
valid_samples = get_input_pipeline(valid_ds,eval_crop_fn,eval_preproc_fn).prefetch(tf.data.experimental.AUTOTUNE)
test_samples = get_input_pipeline(test_ds,eval_preproc_fn).prefetch(tf.data.experimental.AUTOTUNE)

培训和分析

我使用tensorboard profiler来检查第二个批处理大小，并收到一条警告，指出这是高输入限制的，大约40％的处理浪费在输入上。

对于经典的resnet18型号，您可以将批量大小提高到768，而不会出现OOM错误，这就是我所使用的。用bs 256进行的单步操作大约需要2-3秒。

我还收到警告，on_train_batch_size_end与1s批处理时间相比，速度很慢，大约为1.5秒。

模型训练代码是非常简单的keras：

model.fit(
        train_samples,validation_data=valid_samples,epochs=100,batch_size=BS,use_multiprocessing=True
        callbacks=[tensorboard_callback,model_checkpoint_callback,early_stop_callback,reduce_lr_callback]
)

，并将回调指定为：

log_dir = os.path.join(os.getcwd(),'logs')
tensorboard_callback = TensorBoard(log_dir=log_dir,update_freq="epoch",profile_batch=2)

reduce_lr_callback = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss',factor=0.1,patience=5,min_lr=0.001,verbose=1)

model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath='model.{epoch:02d}-{val_loss:.4f}.h5',monitor='val_loss',verbose=1,save_best_only=True)

early_stop_callback = keras.callbacks.EarlyStopping(monitor='val_loss',patience=15)

最后，这是一些张量板配置屏幕快照示例。我不知道如何使它运行得更快：