def yolo_eval(dataset_path, ckpt_path):
"""Yolov3 evaluation."""
ds = create_yolo_dataset(dataset_path, is_training=False)
config = ConfigYOLOV3ResNet18()
net = yolov3_resnet18(config)
eval_net = YoloWithEval(net, config)
print("Load Checkpoint!")
param_dict = load_checkpoint(ckpt_path)
load_param_into_net(net, param_dict)
eval_net.set_train(False)
i = 1.
total = ds.get_dataset_size()
start = time.time()
pred_data = []
print("\n========================================\n")
print("total images num: ", total)
print("Processing, please wait a moment.")
for data in ds.create_dict_iterator():
img_np = data['image']
image_shape = data['image_shape']
annotation = data['annotation']
eval_net.set_train(False)
output = eval_net(Tensor(img_np), Tensor(image_shape))
for batch_idx in range(img_np.shape[0]):
pred_data.append({
"boxes": output[0].asnumpy()[batch_idx],
"box_scores": output[1].asnumpy()[batch_idx],
"annotation": annotation
})
percent = round(i / total * 100, 2)
print(' %s [%d/%d]' % (str(percent) + '%', i, total), end='\r')
i += 1
print(' %s [%d/%d] cost %d ms' %
(str(100.0) + '%', total, total, int((time.time() - start) * 1000)),
end='\n')
precisions, recalls = metrics(pred_data)
print("\n========================================\n")
for i in range(config.num_classes):
print("class {} precision is {:.2f}%, recall is {:.2f}%".format(
i, precisions[i] * 100, recalls[i] * 100))
def main():
parser = argparse.ArgumentParser(description="YOLOv3 train")
parser.add_argument("--only_create_dataset", type=bool, default=False, help="If set it true, only create "
"Mindrecord, default is false.")
parser.add_argument("--distribute", type=bool, default=False, help="Run distribute, default is false.")
parser.add_argument("--device_id", type=int, default=0, help="Device id, default is 0.")
parser.add_argument("--device_num", type=int, default=1, help="Use device nums, default is 1.")
parser.add_argument("--lr", type=float, default=0.001, help="Learning rate, default is 0.001.")
parser.add_argument("--mode", type=str, default="sink", help="Run sink mode or not, default is sink")
parser.add_argument("--epoch_size", type=int, default=10, help="Epoch size, default is 10")
parser.add_argument("--batch_size", type=int, default=32, help="Batch size, default is 32.")
parser.add_argument("--pre_trained", type=str, default=None, help="Pretrained checkpoint file path")
parser.add_argument("--pre_trained_epoch_size", type=int, default=0, help="Pretrained epoch size")
parser.add_argument("--save_checkpoint_epochs", type=int, default=5, help="Save checkpoint epochs, default is 5.")
parser.add_argument("--loss_scale", type=int, default=1024, help="Loss scale, default is 1024.")
parser.add_argument("--mindrecord_dir", type=str, default="./Mindrecord",
help="Mindrecord directory. If the mindrecord_dir is empty, it wil generate mindrecord file by"
"image_dir and anno_path. Note if mindrecord_dir isn't empty, it will use mindrecord_dir "
"rather than image_dir and anno_path. Default is ./Mindrecord_train")
parser.add_argument('--data_url', type=str, default=None, help='Dataset path')
parser.add_argument('--train_url', type=str, default=None, help='Train output path')
parser.add_argument("--anno_path", type=str, default="", help="Annotation path.")
args_opt = parser.parse_args()
device_id = int(os.getenv('DEVICE_ID'))
device_num = int(os.getenv('RANK_SIZE'))
rankid = int(os.getenv('RANK_ID'))
local_data_url = '/cache/data'
local_train_url = '/cache/ckpt'
local_anno_url = '/cache/anno'
local_mindrecord_url = '/cache/mindrecord'
mox.file.copy_parallel(args_opt.mindrecord_dir,local_mindrecord_url)
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=device_id)
if args_opt.distribute:
context.reset_auto_parallel_context()
context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, mirror_mean=True,
device_num=device_num)
init()
rank = rankid
local_train_url = os.path.join(local_train_url,str(device_id))
else:
rank = 0
device_num = 1
print("Start create dataset!")
# It will generate mindrecord file in args_opt.mindrecord_dir,
# and the file name is yolo.mindrecord0, 1, ... file_num.
if not os.path.isdir(local_mindrecord_url):
os.makedirs(local_mindrecord_url)
prefix = "train.mindrecord"
mindrecord_file = os.path.join(local_mindrecord_url, prefix + "0")
if not os.path.exists(mindrecord_file):
mox.file.copy_parallel(args_opt.data_url,local_data_url)
if args_opt.anno_path:
anno_file=os.path.join(local_anno_url,os.path.split(args_opt.anno_path)[1])
mox.file.copy_parallel(args_opt.anno_path,anno_file)
if os.path.isdir(local_data_url) or os.path.exists(anno_file):
print("Create Mindrecord.")
data_to_mindrecord_byte_image(local_data_url,
anno_file,
local_mindrecord_url,
prefix=prefix,
file_num=8)
print("Create Mindrecord Done, at {}".format(args_opt.mindrecord_dir))
mox.file.copy_parallel(local_mindrecord_url,args_opt.mindrecord_dir)
else:
print("image_dir or anno_path not exits.")
if not args_opt.only_create_dataset:
loss_scale = float(args_opt.loss_scale)
# When create MindDataset, using the fitst mindrecord file, such as yolo.mindrecord0.
dataset = create_yolo_dataset(mindrecord_file, repeat_num=args_opt.epoch_size,
batch_size=args_opt.batch_size, device_num=device_num, rank=rank)
dataset_size = dataset.get_dataset_size()
print("Create dataset done!")
net = yolov3_resnet18(ConfigYOLOV3ResNet18())
net = YoloWithLossCell(net, ConfigYOLOV3ResNet18())
init_net_param(net, "XavierUniform")
# checkpoint
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_size * args_opt.save_checkpoint_epochs)
ckpoint_cb = ModelCheckpoint(prefix="yolov3", directory=local_train_url, config=ckpt_config)
if args_opt.pre_trained:
if args_opt.pre_trained_epoch_size <= 0:
raise KeyError("pre_trained_epoch_size must be greater than 0.")
param_dict = load_checkpoint(args_opt.pre_trained)
load_param_into_net(net, param_dict)
total_epoch_size = 60
if args_opt.distribute:
total_epoch_size = 160
lr = Tensor(get_lr(learning_rate=args_opt.lr, start_step=args_opt.pre_trained_epoch_size * dataset_size,
global_step=total_epoch_size * dataset_size,
decay_step=1000, decay_rate=0.95, steps=True))
opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()), lr, loss_scale=loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
callback = [TimeMonitor(data_size=dataset_size), LossMonitor(), ckpoint_cb]
model = Model(net)
dataset_sink_mode = False
if args_opt.mode == "sink":
print("In sink mode, one epoch return a loss.")
dataset_sink_mode = True
print("Start train YOLOv3, the first epoch will be slower because of the graph compilation.")
model.train(args_opt.epoch_size, dataset, callbacks=callback, dataset_sink_mode=dataset_sink_mode)
if device_id ==1:
mox.file.copy_parallel(local_train_url,args_opt.train_url)
def test_yolov3():
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend")
rank = 0
device_num = 1
lr_init = 0.001
epoch_size = 3
batch_size = 32
loss_scale = 1024
mindrecord_dir = DATA_DIR
# It will generate mindrecord file in args_opt.mindrecord_dir,
# and the file name is yolo.mindrecord0, 1, ... file_num.
if not os.path.isdir(mindrecord_dir):
raise KeyError("mindrecord path is not exist.")
prefix = "yolo.mindrecord"
mindrecord_file = os.path.join(mindrecord_dir, prefix + "0")
print("yolov3 mindrecord is ", mindrecord_file)
if not os.path.exists(mindrecord_file):
print("mindrecord file is not exist.")
assert False
else:
loss_scale = float(loss_scale)
# When create MindDataset, using the fitst mindrecord file, such as yolo.mindrecord0.
dataset = create_yolo_dataset(mindrecord_file,
repeat_num=1,
batch_size=batch_size,
device_num=device_num,
rank=rank)
dataset_size = dataset.get_dataset_size()
print("Create dataset done!")
net = yolov3_resnet18(ConfigYOLOV3ResNet18())
net = YoloWithLossCell(net, ConfigYOLOV3ResNet18())
total_epoch_size = 60
lr = Tensor(
get_lr(learning_rate=lr_init,
start_step=0,
global_step=total_epoch_size * dataset_size,
decay_step=1000,
decay_rate=0.95,
steps=True))
opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()),
lr,
loss_scale=loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
model_callback = ModelCallback()
time_monitor_callback = TimeMonitor(data_size=dataset_size)
callback = [model_callback, time_monitor_callback]
model = Model(net)
print(
"Start train YOLOv3, the first epoch will be slower because of the graph compilation."
)
model.train(epoch_size,
dataset,
callbacks=callback,
dataset_sink_mode=True,
sink_size=dataset.get_dataset_size())
# assertion occurs while the loss value, overflow state or loss_scale value is wrong
loss_value = np.array(model_callback.loss_list)
expect_loss_value = [6600, 4200, 2700]
print("loss value: {}".format(loss_value))
assert loss_value[0] < expect_loss_value[0]
assert loss_value[1] < expect_loss_value[1]
assert loss_value[2] < expect_loss_value[2]
epoch_mseconds = np.array(time_monitor_callback.epoch_mseconds_list)[2]
expect_epoch_mseconds = 950
print("epoch mseconds: {}".format(epoch_mseconds))
assert epoch_mseconds <= expect_epoch_mseconds
per_step_mseconds = np.array(
time_monitor_callback.per_step_mseconds_list)[2]
expect_per_step_mseconds = 110
print("per step mseconds: {}".format(per_step_mseconds))
assert per_step_mseconds <= expect_per_step_mseconds
print("yolov3 test case passed.")
def main():
parser = argparse.ArgumentParser(description="YOLOv3 train")
parser.add_argument(
"--only_create_dataset",
type=ast.literal_eval,
default=False,
help="If set it true, only create Mindrecord, default is False.")
parser.add_argument("--distribute",
type=ast.literal_eval,
default=False,
help="Run distribute, default is False.")
parser.add_argument("--device_id",
type=int,
default=0,
help="Device id, default is 0.")
parser.add_argument("--device_num",
type=int,
default=1,
help="Use device nums, default is 1.")
parser.add_argument("--lr",
type=float,
default=0.001,
help="Learning rate, default is 0.001.")
parser.add_argument("--mode",
type=str,
default="sink",
help="Run sink mode or not, default is sink")
parser.add_argument("--epoch_size",
type=int,
default=50,
help="Epoch size, default is 50")
parser.add_argument("--batch_size",
type=int,
default=32,
help="Batch size, default is 32.")
parser.add_argument("--pre_trained",
type=str,
default=None,
help="Pretrained checkpoint file path")
parser.add_argument("--pre_trained_epoch_size",
type=int,
default=0,
help="Pretrained epoch size")
parser.add_argument("--save_checkpoint_epochs",
type=int,
default=5,
help="Save checkpoint epochs, default is 5.")
parser.add_argument("--loss_scale",
type=int,
default=1024,
help="Loss scale, default is 1024.")
parser.add_argument(
"--mindrecord_dir",
type=str,
default="./Mindrecord_train",
help=
"Mindrecord directory. If the mindrecord_dir is empty, it wil generate mindrecord file by "
"image_dir and anno_path. Note if mindrecord_dir isn't empty, it will use mindrecord_dir "
"rather than image_dir and anno_path. Default is ./Mindrecord_train")
parser.add_argument("--image_dir",
type=str,
default="",
help="Dataset directory, "
"the absolute image path is joined by the image_dir "
"and the relative path in anno_path")
parser.add_argument("--anno_path",
type=str,
default="",
help="Annotation path.")
args_opt = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
device_id=args_opt.device_id)
if args_opt.distribute:
device_num = args_opt.device_num
context.reset_auto_parallel_context()
context.set_auto_parallel_context(
parallel_mode=ParallelMode.DATA_PARALLEL,
gradients_mean=True,
device_num=device_num)
init()
rank = args_opt.device_id % device_num
else:
rank = 0
device_num = 1
print("Start create dataset!")
# It will generate mindrecord file in args_opt.mindrecord_dir,
# and the file name is yolo.mindrecord0, 1, ... file_num.
if not os.path.isdir(args_opt.mindrecord_dir):
os.makedirs(args_opt.mindrecord_dir)
prefix = "yolo.mindrecord"
mindrecord_file = os.path.join(args_opt.mindrecord_dir, prefix + "0")
if not os.path.exists(mindrecord_file):
if os.path.isdir(args_opt.image_dir) and os.path.exists(
args_opt.anno_path):
print("Create Mindrecord.")
data_to_mindrecord_byte_image(args_opt.image_dir,
args_opt.anno_path,
args_opt.mindrecord_dir, prefix, 8)
print("Create Mindrecord Done, at {}".format(
args_opt.mindrecord_dir))
else:
raise ValueError('image_dir {} or anno_path {} does not exist'.format(\
args_opt.image_dir, args_opt.anno_path))
if not args_opt.only_create_dataset:
loss_scale = float(args_opt.loss_scale)
# When create MindDataset, using the fitst mindrecord file, such as yolo.mindrecord0.
dataset = create_yolo_dataset(mindrecord_file,
batch_size=args_opt.batch_size,
device_num=device_num,
rank=rank)
dataset_size = dataset.get_dataset_size()
print("Create dataset done!")
net = yolov3_resnet18(ConfigYOLOV3ResNet18())
net = YoloWithLossCell(net, ConfigYOLOV3ResNet18())
init_net_param(net, "XavierUniform")
# checkpoint
ckpt_config = CheckpointConfig(save_checkpoint_steps=dataset_size *
args_opt.save_checkpoint_epochs)
ckpoint_cb = ModelCheckpoint(prefix="yolov3",
directory=None,
config=ckpt_config)
if args_opt.pre_trained:
if args_opt.pre_trained_epoch_size <= 0:
raise KeyError(
"pre_trained_epoch_size must be greater than 0.")
param_dict = load_checkpoint(args_opt.pre_trained)
load_param_into_net(net, param_dict)
total_epoch_size = 60
if args_opt.distribute:
total_epoch_size = 160
lr = Tensor(
get_lr(learning_rate=args_opt.lr,
start_step=args_opt.pre_trained_epoch_size * dataset_size,
global_step=total_epoch_size * dataset_size,
decay_step=1000,
decay_rate=0.95,
steps=True))
opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()),
lr,
loss_scale=loss_scale)
net = TrainingWrapper(net, opt, loss_scale)
callback = [
TimeMonitor(data_size=dataset_size),
LossMonitor(), ckpoint_cb
]
model = Model(net)
dataset_sink_mode = False
if args_opt.mode == "sink":
print("In sink mode, one epoch return a loss.")
dataset_sink_mode = True
print(
"Start train YOLOv3, the first epoch will be slower because of the graph compilation."
)
model.train(args_opt.epoch_size,
dataset,
callbacks=callback,
dataset_sink_mode=dataset_sink_mode)
type=str,
default="yolov3_resnet18.air",
help="output file name.")
parser.add_argument('--file_format',
type=str,
choices=["AIR", "ONNX", "MINDIR"],
default='AIR',
help='file format')
args = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target="Ascend",
device_id=args.device_id)
if __name__ == "__main__":
config = ConfigYOLOV3ResNet18()
network = yolov3_resnet18(config)
param_dict = load_checkpoint(args.ckpt_file)
load_param_into_net(network, param_dict)
network.set_train(False)
shape = [args.batch_size, 3] + config.img_shape
input_data = Tensor(np.zeros(shape), ms.float32)
export(network,
input_data,
file_name=args.file_name,
file_format=args.file_format)
def metrics(pred_data):
"""Calculate precision and recall of predicted bboxes."""
config = ConfigYOLOV3ResNet18()
num_classes = config.num_classes
count_corrects = [1e-6 for _ in range(num_classes)]
count_grounds = [1e-6 for _ in range(num_classes)]
count_preds = [1e-6 for _ in range(num_classes)]
for i, sample in enumerate(pred_data):
gt_anno = sample["annotation"]
box_scores = sample['box_scores']
boxes = sample['boxes']
mask = box_scores >= config.obj_threshold
boxes_ = []
scores_ = []
classes_ = []
max_boxes = config.nms_max_num
for c in range(num_classes):
class_boxes = np.reshape(boxes,
[-1, 4])[np.reshape(mask[:, c], [-1])]
class_box_scores = np.reshape(box_scores[:, c],
[-1])[np.reshape(mask[:, c], [-1])]
nms_index = apply_nms(class_boxes, class_box_scores,
config.nms_threshold, max_boxes)
class_boxes = class_boxes[nms_index]
class_box_scores = class_box_scores[nms_index]
classes = np.ones_like(class_box_scores, 'int32') * c
boxes_.append(class_boxes)
scores_.append(class_box_scores)
classes_.append(classes)
boxes = np.concatenate(boxes_, axis=0)
classes = np.concatenate(classes_, axis=0)
# metric
count_correct = [1e-6 for _ in range(num_classes)]
count_ground = [1e-6 for _ in range(num_classes)]
count_pred = [1e-6 for _ in range(num_classes)]
for anno in gt_anno:
count_ground[anno[4]] += 1
for box_index, box in enumerate(boxes):
bbox_pred = [box[1], box[0], box[3], box[2]]
count_pred[classes[box_index]] += 1
for anno in gt_anno:
class_ground = anno[4]
if classes[box_index] == class_ground:
iou = calc_iou(bbox_pred, anno)
if iou >= 0.5:
count_correct[class_ground] += 1
break
count_corrects = [
count_corrects[i] + count_correct[i] for i in range(num_classes)
]
count_preds = [
count_preds[i] + count_pred[i] for i in range(num_classes)
]
count_grounds = [
count_grounds[i] + count_ground[i] for i in range(num_classes)
]
precision = np.array(
[count_corrects[ix] / count_preds[ix] for ix in range(num_classes)])
recall = np.array(
[count_corrects[ix] / count_grounds[ix] for ix in range(num_classes)])
return precision, recall