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: trueclip_boxes: falsevariances: "[0.1, 0.1, 0.2, 0.2]"arch: "resnet"nlayers: 18freeze_bn: falsefreeze_blocks: 0
}
training_config {batch_size_per_gpu: 32num_epochs: 80enable_qat: falselearning_rate {soft_start_annealing_schedule {min_learning_rate: 5e-5max_learning_rate: 2e-2soft_start: 0.15annealing: 0.8}}regularizer {type: L1weight: 3e-5}
}
eval_config {validation_period_during_training: 5average_precision_mode: SAMPLEbatch_size: 16matching_iou_threshold: 0.5
}
nms_config {confidence_threshold: 0.01clustering_iou_threshold: 0.6top_k: 200
}
augmentation_config {output_width: 300output_height: 300output_channel: 3
}
dataset_config {data_sources: {tfrecords_path: "/home/alex/7th_sky_hackathon/data/tfrecords/kitti_train*"}include_difficult_in_training: truetarget_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

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