TensorFlow目标检测API训练Centernet时CPU/RAM占用高GPU无占用问题

AI 小北

2026-10-11

TensorFlow Object Detection API训练CenterNet时GPU无调用、CPU内存占用过高问题

问题现象

使用TensorFlow Object Detection API训练模型时，CPU和内存占用极高，Windows任务管理器显示GPU基本未被调用：
任务管理器资源占用截图

环境验证情况

已完成TF目标检测API安装，且验证GPU可被TensorFlow正常识别，验证命令及输出如下：

python -c "import tensorflow as tf;print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

输出：

2021-07-20 15:36:37.630320: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cudart64_110.dll
2021-07-20 15:36:49.683811: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library nvcuda.dll
2021-07-20 15:36:49.990907: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties:
pciBusID: 0000:01:00.0 name: NVIDIA GeForce RTX 3080 Laptop GPU computeCapability: 8.6
coreClock: 1.605GHz coreCount: 48 deviceMemorySize: 8.00GiB deviceMemoryBandwidth: 417.29GiB/s
2021-07-20 15:36:50.017685: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cudart64_110.dll
2021-07-20 15:36:50.142257: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cublas64_11.dll
2021-07-20 15:36:50.158525: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cublasLt64_11.dll
2021-07-20 15:36:50.173970: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cufft64_10.dll
2021-07-20 15:36:50.183516: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library curand64_10.dll
2021-07-20 15:36:50.196516: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cusolver64_11.dll
2021-07-20 15:36:50.213625: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cusparse64_11.dll
2021-07-20 15:36:50.231417: I tensorflow/stream_executor/platform/default/dso_loader.cc:53] Successfully opened dynamic library cudnn64_8.dll
2021-07-20 15:36:50.234253: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
2021-07-20 15:36:50.238133: I tensorflow/core/platform/cpu_feature_guard.cc:142] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX AVX2
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2021-07-20 15:36:50.245602: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1733] Found device 0 with properties:
pciBusID: 0000:01:00.0 name: NVIDIA GeForce RTX 3080 Laptop GPU computeCapability: 8.6
coreClock: 1.605GHz coreCount: 48 deviceMemorySize: 8.00GiB deviceMemoryBandwidth: 417.29GiB/s
2021-07-20 15:36:50.265550: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1871] Adding visible gpu devices: 0
2021-07-20 15:36:54.162700: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1258] Device interconnect StreamExecutor with strength 1 edge matrix:
2021-07-20 15:36:54.168506: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1264]      0
2021-07-20 15:36:54.169910: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1277] 0:   N
2021-07-20 15:36:54.176538: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1418] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 5177 MB memory) -> physical GPU (device: 0, name: NVIDIA GeForce RTX 3080 Laptop GPU, pci bus id: 0000:01:00.0, compute capability: 8.6)
tf.Tensor(1222.1837, shape=(), dtype=float32)

补充说明

该问题仅在训练centernet_hg104_512x512_coco17_tpu-8模型时出现，训练ssd_resnet、efficientdet等其他模型时GPU可正常调用，本次使用的pipeline.config配置如下：

model {
  center_net {
    num_classes: 1
    feature_extractor {
      type: "hourglass_104"
      channel_means: 104.01361846923828
      channel_means: 114.03422546386719
      channel_means: 119.91659545898438
      channel_stds: 73.60276794433594
      channel_stds: 69.89082336425781
      channel_stds: 70.91507720947266
      bgr_ordering: true
    }
    image_resizer {
      keep_aspect_ratio_resizer {
        min_dimension: 512
        max_dimension: 512
        pad_to_max_dimension: true
      }
    }
    object_detection_task {
      task_loss_weight: 1.0
      offset_loss_weight: 1.0
      scale_loss_weight: 0.10000000149011612
      localization_loss {
        l1_localization_loss {
        }
      }
    }
    object_center_params {
      object_center_loss_weight: 1.0
      classification_loss {
        penalty_reduced_logistic_focal_loss {
          alpha: 2.0
          beta: 4.0
        }
      }
      min_box_overlap_iou: 0.6
      max_box_predictions: 50
    }
  }
}
train_config {
  batch_size: 2
  data_augmentation_options {
    random_horizontal_flip {
    }
  }
  data_augmentation_options {
    random_crop_image {
      min_aspect_ratio: 0.5
      max_aspect_ratio: 1.7000000476837158
      random_coef: 0.25
    }
  }
  data_augmentation_options {
    random_adjust_hue {
    }
  }
  data_augmentation_options {
    random_adjust_contrast {
    }
  }
  data_augmentation_options {
    random_adjust_saturation {
    }
  }
  data_augmentation_options {
    random_adjust_brightness {
    }
  }
  data_augmentation_options {
    random_absolute_pad_image {
      max_height_padding: 200
      max_width_padding: 200
      pad_color: 0.0
      pad_color: 0.0
      pad_color: 0.0
    }
  }
  optimizer {
    adam_optimizer {
      learning_rate {
        manual_step_learning_rate {
          initial_learning_rate: 0.0010000000474974513
          schedule {
            step: 1000
            learning_rate: 9.999999747378752e-05
          }
          schedule {
            step: 5000
            learning_rate: 9.999999747378752e-06
          }
        }
      }
      epsilon: 1.0000000116860974e-07
    }
    use_moving_average: false
  }
  fine_tune_checkpoint: "pre-trained-models/centernet_hg104_512x512_coco17_tpu-8/checkpoint/ckpt-0"
  num_steps: 5000
  max_number_of_boxes: 50
  unpad_groundtruth_tensors: false
  fine_tune_checkpoint_type: "detection"
  fine_tune_checkpoint_version: V2
}
train_input_reader {
  label_map_path: "annotations/label_map.pbtxt"
  tf_record_input_reader {
    input_path: "annotations/train.record"
  }
}
eval_config {
  metrics_set: "coco_detection_metrics"
  use_moving_averages: false
  batch_size: 1
}
eval_input_reader {
  label_map_path: "annotations/label_map.pbtxt"
  shuffle: false
  num_epochs: 1
  tf_record_input_reader {
    input_path: "annotations/test.record"
  }
}

当前batch_size设置为2，若调大batch_size，训练脚本会卡顿一段时间后退出。

软硬件环境

操作系统：Windows 10
CPU：i9-10980HK
内存：32GB
GPU：GTX3080 8GB专用显存
TensorFlow版本：2.5
CUDA版本：11.3.1
cuDNN版本：8.2.1.32

问题解答

这个现象不属于预期，核心原因和CenterNet的Hourglass骨干网络以及TFODAPI的默认配置有关，可按以下步骤排查调整：

排查数据加载瓶颈
你配置的7种数据增强操作默认全部在CPU侧执行，Hourglass网络本身计算量极大，很容易出现CPU预处理速度跟不上GPU计算速度的情况，导致GPU空转。可以先临时注释掉所有数据增强选项，观察GPU占用是否有明显上升，如果有，就说明是预处理瓶颈，可以在train_input_reader中添加num_readers: 4、prefetch_size: 8参数，提升数据加载并行度。
修复TPU预训练权重的适配问题
你使用的centernet_hg104_512x512_coco17_tpu-8是官方为TPU训练优化的预训练权重，部分算子在GPU上没有对应的高效实现，会自动回退到CPU执行。可以尝试更换为非TPU版本的CenterNet预训练权重，或者在pipeline.config的train_config块中添加use_bfloat16: false配置，强制启用GPU算子调度。
解决显存不足导致的算子回退
8GB显存运行batch_size=2的HG104模型本身已经接近显存阈值，TF会自动把部分大显存占用的算子调度到CPU执行，这也是你调大batch_size就崩溃的原因。可以尝试开启显存动态分配，在训练脚本最开头添加以下代码：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
  except RuntimeError as e:
    print(e)

如果还是不行，建议更换为轻量版的CenterNet骨干，比如hg52，或者换用mobilenet骨干的CenterNet模型，适配8GB显存的设备。
4. 修复版本兼容性问题
TensorFlow 2.5对30系显卡的安培架构适配存在部分缺陷，你使用的CUDA11.3也比TF2.5官方支持的CUDA11.2版本更高，部分算子适配异常也会导致CPU回退。可以尝试升级到TensorFlow 2.6+版本，匹配对应的CUDA/cuDNN版本，解决架构适配问题。

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

TensorFlow目标检测API训练Centernet时CPU/RAM占用高GPU无占用问题

问题现象 #

环境验证情况 #

补充说明 #

软硬件环境 #

问题解答 #

问题现象

环境验证情况

补充说明

软硬件环境

问题解答