diff --git a/model_zoo/official/cv/maskrcnn/README.md b/model_zoo/official/cv/maskrcnn/README.md index 89a123b6ca..fadce5ad32 100644 --- a/model_zoo/official/cv/maskrcnn/README.md +++ b/model_zoo/official/cv/maskrcnn/README.md @@ -35,7 +35,7 @@ MaskRcnn is a two-stage target detection network,This network uses a region prop └─train2017 ``` - + Notice that the coco2017 dataset will be converted to MindRecord which is a data format in MindSpore. The dataset conversion may take about 4 hours. 2. If your own dataset is used. **Select dataset to other when run script.** Organize the dataset infomation into a TXT file, each row in the file is as follows: diff --git a/model_zoo/official/cv/maskrcnn/src/config.py b/model_zoo/official/cv/maskrcnn/src/config.py index 38c43e3b2f..97028c782e 100644 --- a/model_zoo/official/cv/maskrcnn/src/config.py +++ b/model_zoo/official/cv/maskrcnn/src/config.py @@ -134,6 +134,7 @@ config = ed({ "loss_scale": 1, "momentum": 0.91, "weight_decay": 1e-4, + "pretrain_epoch_size": 0, "epoch_size": 12, "save_checkpoint": True, "save_checkpoint_epochs": 1, diff --git a/model_zoo/official/cv/maskrcnn/src/lr_schedule.py b/model_zoo/official/cv/maskrcnn/src/lr_schedule.py index 29779c9184..dae7e2ea82 100644 --- a/model_zoo/official/cv/maskrcnn/src/lr_schedule.py +++ b/model_zoo/official/cv/maskrcnn/src/lr_schedule.py @@ -25,7 +25,7 @@ def a_cosine_learning_rate(current_step, base_lr, warmup_steps, decay_steps): learning_rate = (1 + math.cos(base * math.pi)) / 2 * base_lr return learning_rate -def dynamic_lr(config, rank_size=1): +def dynamic_lr(config, rank_size=1, start_steps=0): """dynamic learning rate generator""" base_lr = config.base_lr @@ -38,5 +38,5 @@ def dynamic_lr(config, rank_size=1): lr.append(linear_warmup_learning_rate(i, warmup_steps, base_lr, base_lr * config.warmup_ratio)) else: lr.append(a_cosine_learning_rate(i, base_lr, warmup_steps, total_steps)) - - return lr + learning_rate = lr[start_steps:] + return learning_rate diff --git a/model_zoo/official/cv/maskrcnn/train.py b/model_zoo/official/cv/maskrcnn/train.py index 27936b828e..8df2fab00e 100644 --- a/model_zoo/official/cv/maskrcnn/train.py +++ b/model_zoo/official/cv/maskrcnn/train.py @@ -108,13 +108,15 @@ if __name__ == '__main__': load_path = args_opt.pre_trained if load_path != "": param_dict = load_checkpoint(load_path) - for item in list(param_dict.keys()): - if not (item.startswith('backbone') or item.startswith('rcnn_mask')): - param_dict.pop(item) + if config.pretrain_epoch_size == 0: + for item in list(param_dict.keys()): + if not (item.startswith('backbone') or item.startswith('rcnn_mask')): + param_dict.pop(item) load_param_into_net(net, param_dict) loss = LossNet() - lr = Tensor(dynamic_lr(config, rank_size=device_num), mstype.float32) + lr = Tensor(dynamic_lr(config, rank_size=device_num, start_steps=config.pretrain_epoch_size * dataset_size), + mstype.float32) opt = SGD(params=net.trainable_params(), learning_rate=lr, momentum=config.momentum, weight_decay=config.weight_decay, loss_scale=config.loss_scale)