Запутаться в отношении data_format в tf.layer.cond2d
Учитывая, что документация на https://www.tensorflow.org/api_docs/python/tf/layers/conv2d упоминает, что параметр data_format 'channel_last', который я предполагаю быть таким же, как: NHWC, но при запуске модели я сталкиваюсь с ошибкой нехватки ресурсов памяти, которая показывает мне следующее:
ResourceExhaustedError (see above for traceback): OOM when allocating tensor with shape[510,256,28,28] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
[[Node: block_2/conv_1/Conv2D = Conv2D[T=DT_FLOAT, data_format="NCHW", dilations=[1, 1, 1, 1], padding="SAME", strides=[1, 1, 1, 1], use_cudnn_on_gpu=true, _device="/job:localhost/replica:0/task:0/device:GPU:0"](block_2/relu_0, block_2/conv_1/kernel/read)]]
И в указанной выше ошибке я вижу, что используется data_format: "NCHW", Так что я хотел бы знать правду data_format используемый.
Я спрашиваю, так как я знаю, что это может иметь некоторые дополнительные характеристики скорости.
Обратите внимание, что я использую tenorflow.
Кроме того, я использую tf.data.Dataset API в тензорном потоке. Вот код для получения tfrecords:
dataset_train = tf.data.TFRecordDataset(file_name).map(_parse_function).batch(train_batch_size)
Любая помощь высоко ценится!!