NVIDIA 7th SkyHackathon（六）Tao 目标检测模型训练与评估

1.模型准备

1.1 安装 NGC CLI

使用 NGC CLI 来获得预训练模型，关于 NGC 的详细信息可访问 官网 中的 SETUP 了解

# 创建 NGC CLI 目录
%env CLI=ngccli_cat_linux.zip
mkdir -p $LOCAL_PROJECT_DIR/ngccli

# 移除之前安装的 NGC CLI
rm -rf $LOCAL_PROJECT_DIR/ngccli/*

# 下载并解压 NGC CLI
wget "https://ngc.nvidia.com/downloads/$CLI" -P $LOCAL_PROJECT_DIR/ngccli
unzip -u "$LOCAL_PROJECT_DIR/ngccli/$CLI" -d $LOCAL_PROJECT_DIR/ngccli/
rm $LOCAL_PROJECT_DIR/ngccli/*.zip 

# 添加 NGC CLI 到环境变量中
os.environ["PATH"]="{}/ngccli:{}".format(os.getenv("LOCAL_PROJECT_DIR", ""), os.getenv("PATH", ""))

1.2 安装预训练模型

# 查看目前 NGC 上可下载的关于目标检测的预训练模型
ngc registry model list nvidia/tao/pretrained_object_detection:*

# 创建保存预训练模型的目录
mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/

# 下载选择的预训练模型
ngc registry model download-version nvidia/tao/pretrained_object_detection:resnet18 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18

#查看预训练模型是否下载成功
ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/pretrained_object_detection_vresnet18

1.3 设置模型参数

打开 $LOCAL_SPECS_DIR/ssd_train_resnet18_kitti.txt，根据实际需要，修改模型参数，其中：

• batch_size_per_gpu：定义 batch_size 大小，（若 GPU 显存不大，可以调小一点，以免出现 out of memory）
• num_epochs：定义会训练多少轮，第一次训练建议不小于 80
• validation_period_during_training：定义训练验证轮次，可直观看出训练的效果变化

random_seed: 42
ssd_config {
  aspect_ratios_global: "[1.0, 2.0, 0.5, 3.0, 1.0/3.0]"
  scales: "[0.05, 0.1, 0.25, 0.4, 0.55, 0.7, 0.85]"
  two_boxes_for_ar1: true
  clip_boxes: false
  variances: "[0.1, 0.1, 0.2, 0.2]"
  arch: "resnet"
  nlayers: 18
  freeze_bn: false
  freeze_blocks: 0
}
training_config {
  batch_size_per_gpu: 32
  num_epochs: 80
  enable_qat: false
  learning_rate {
  soft_start_annealing_schedule {
    min_learning_rate: 5e-5
    max_learning_rate: 2e-2
    soft_start: 0.15
    annealing: 0.8
    }
  }
  regularizer {
    type: L1
    weight: 3e-5
  }
}
eval_config {
  validation_period_during_training: 5
  average_precision_mode: SAMPLE
  batch_size: 16
  matching_iou_threshold: 0.5
}
nms_config {
  confidence_threshold: 0.01
  clustering_iou_threshold: 0.6
  top_k: 200
}
augmentation_config {
    output_width: 300
    output_height: 300
    output_channel: 3
}
dataset_config {
  data_sources: {
    tfrecords_path: "/home/alex/7th_sky_hackathon/data/tfrecords/kitti_train*"
  }
  include_difficult_in_training: true
  target_class_mapping {
      key: "cardboard"
      value: "cardboard"
  }
  target_class_mapping {
      key: "bottle"
      value: "bottle"
  }
  target_class_mapping {
      key: "banane"
      value: "banane"
  }
  validation_data_sources: {
      label_directory_path: "/home/alex/7th_sky_hackathon/data/val/label"
      image_directory_path: "/home/alex/7th_sky_hackathon/data/val/image"
  }
}

2.模型训练

首先在 $LOCAL_EXPERIMENT_DIR 中创建模型输出文件夹 experiment_dir_unpruned_final，保存训练结果

之后，使用 tao 套件进行训练，各参数如下：

• –gpus：几个 GPU
• –gup_index：使用哪个 GPU
• -e：训练的设置文件
• -r：训练模型输出文件夹
• -k：秘钥
• -m：预训练模型

tao ssd train --gpus 1 -- gpu_index=$GPU_INDEX -e $SPECS_DIR/ssd_train_resnet18_kitti.txt -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final -k $KEY -m $USER_EXPERIMENT_DIR/7th/pretrained_resnet18/pretrained_object_detection_vresnet18/resnet_18.hdf5

训练完成后，可在 $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final/weights 目录下，查看每轮训练出的模型

文件 ssd_training_log_resnet18.csv 中，记录了所有模型的精度，文件具体内容如下：

epoch,AP_banane,AP_bottle,AP_cardboard,loss,lr,mAP,validation_loss
1,nan,nan,nan,27.146854,8.2377446e-05,nan,nan
2,nan,nan,nan,17.637777,0.00013572087,nan,nan
3,nan,nan,nan,15.214624,0.00022360678,nan,nan
4,nan,nan,nan,13.826463,0.00036840313,nan,nan

选择 mAP 最大的模型作为最好的模型

%set_env EPOCH=080

3.模型评估

选择好模型后，使用 tao 套件，可对模型进行评估

tao ssd evaluate --gpu_index=$GPU_INDEX  -e $SPECS_DIR/ssd_train_resnet18_kitti.txt -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final/weights/ssd_resnet18_epoch_$EPOCH.tlt -k $KEY

4.模型剪枝

首先在目录 $USER_EXPERIMENT_DIR 下建立存储剪枝模型的目录 experiment_dir_pruned

之后使用 tao 工具套件对模型进行剪枝，各参数如下：

• -m：要剪枝的模型
• -o：剪枝后的输出
• -ep：剪枝的方式
• -pth：剪枝的阈值
• -k：秘钥

tao ssd prune --gpu_index=$GPU_INDEX -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned_final/weights/ssd_resnet18_epoch_$EPOCH.tlt -o $USER_EXPERIMENT_DIR/experiment_dir_pruned/ssd_resnet18_pruned.tlt -eq intersection -pth 0.6 -k $KEY

在经过剪枝后，需要对模型重新训练、评估，步骤同上