|
-
- import os
- import argparse
- import ast
- import math
- import mindspore.nn as nn
- import mindspore.ops
- from mindspore import context, Tensor
- import numpy as np
- from mindspore.communication.management import init, get_rank
- from mindspore.train.callback import CheckpointConfig, ModelCheckpoint, LossMonitor, TimeMonitor, Callback
- from mindspore.train import Model
- from mindspore.context import ParallelMode
- from mindspore.train.serialization import load_checkpoint, load_param_into_net
- from mindspore.common import set_seed
- from src.config import config
- from src.dataset import create_EfficientDet_datasets
- import mindspore as ms
- from src.lr_schedule import get_lr_cosine
- # from src.init_params import init_net_param, filter_checkpoint_parameter
- from src.mind_backbone import EfficientDetBackbone
- # from src.efficientdet.mind_loss import FocalLoss
- from src.efficientdet.loss import FocalLoss
- # from loss import FocalLoss
- from mindspore.nn import TrainOneStepCell
- import mindspore.common.initializer as weight_init
-
- from mindspore.common import initializer as initier
- from src.efficientnet.model import EfficientNet
- from src.efficientdet.model import BiFPN,Classifier,Regressor
- # 算子性能分析
- from src.monitor import Monitor
- from mindspore.profiler import Profiler
-
- # 云上训练代码的改动之处
- # import moxing as mox # 将数据拷贝到脚本的包
- set_seed(0)
-
- def init_weights(model):
- # 返回所有模块的迭代器
- for name, cell in model.cells_and_names():
- is_conv_layer = isinstance(cell, nn.Conv2d)
-
- if is_conv_layer:
- cell.weight.set_data(initier.initializer(initier.Constant(0.6), cell.weight.shape))
- if cell.has_bias is True:
- cell.bias.set_data(initier.initializer('zeros', cell.bias.shape))
-
- class WithLossCell(nn.Cell):
- def __init__(self, backbone, loss):
- super(WithLossCell, self).__init__()
- self.backbone = backbone
- self.loss = loss
- self.print_op = mindspore.ops.Print()
-
- def construct(self, x, y):
- # self.print_op("img_id :")
-
- _, reg, cls, anchor = self.backbone(x)
-
- cls_loss, reg_loss = self.loss(reg, cls, anchor, y)
-
- return cls_loss + reg_loss
-
-
- def main():
-
- parser = argparse.ArgumentParser(description="EfficientDet training")
- parser.add_argument("--distribute", type=ast.literal_eval, default=False, help="Run distribute, default is False.")
- parser.add_argument("--workers", type=int, default=24, help="Num parallel workers.")
- parser.add_argument("--data_url", type=str, default="/home/work/user-job-dir/inputs/data/", help="mindrecord dir")
- parser.add_argument("--train_url", type=str, default="/home/work/user-job-dir/model/", help="ckpt output dir in obs")
- parser.add_argument("--lr", type=float, default=0.05, help="Learning rate, default is 0.1.")
- 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=5, help="Epoch size, default is 500.")
- parser.add_argument("--batch_size", type=int, default=8, help="Batch size, default is 32.")
- parser.add_argument("--pre_trained", type=str, default="/data/efficientdet_ch/efdet.ckpt", help="Pretrained Checkpoint file path.")
- parser.add_argument("--pre_trained_epoch_size", type=int, default=0, help="Pretrained epoch size.")
- parser.add_argument("--pretrained_backbone", type=str, default=None, help="backbone ckpt file path.")
- parser.add_argument("--save_checkpoint_epochs", type=int, default=1, 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("--filter_weight", type=ast.literal_eval, default=False, help="Filter weight parameters, default is False.")
- parser.add_argument("--run_platform", type=str, default="Ascend", choices="Ascend", help="run platform, only support Ascend.")
-
- args_opt = parser.parse_args()
-
- device_id = int(os.getenv("DEVICE_ID"))
-
- if args_opt.run_platform == "Ascend":
- # context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", enable_reduce_precision=True) # save_graphs=True
- context.set_context(mode=context.PYNATIVE_MODE, device_target="Ascend") # save_graphs=True
- if args_opt.distribute:
- if os.getenv("DEVICE_ID", "not_set").isdigit():
- context.set_context(device_id=int(os.getenv("DEVICE_ID")))
- init()
- device_num = int(os.getenv("DEVICE_NUM"))
- rank = int(os.getenv("RANK_ID"))
- rank_size = int(os.getenv("RANK_SIZE"))
- context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, gradients_mean=True,
- device_num=device_num)
- else:
- rank = 0
- device_num = 1
- context.set_context(device_id=device_id, save_graphs=True)
-
- else:
- raise ValueError("Unsupported platform.")
- # 生成EfficientDet.mindrecord
- mindrecord_file = os.path.join(config.mindrecord_dir, "EfficientDet.mindrecord0")
-
- dataset = create_EfficientDet_datasets(mindrecord_file, repeat_num=1,
- num_parallel_workers=args_opt.workers,
- batch_size=args_opt.batch_size, device_num=device_num, rank=rank)
- dataset_size = dataset.get_dataset_size()
-
- print("Create dataset done!")
-
- # profiler = Profiler(output_path='./profile2', is_detail=True, is_show_op_path=False)
- net = EfficientDetBackbone(90, 0, False, False) # 屏蔽掉随机
-
- if args_opt.pre_trained:
- param_dict = load_checkpoint(args_opt.pre_trained)
- load_param_into_net(net, param_dict)
-
- init_weights(net)
-
- loss = FocalLoss()
- net = WithLossCell(net, loss)
-
- net.set_train()
- loss_scale = float(args_opt.loss_scale)
- # 此处的传入参数是none
-
- lr = Tensor(get_lr_cosine(init_lr=0.01, steps_per_epoch=dataset_size, warmup_epochs=int(args_opt.epoch_size / 15),
- max_epoch=args_opt.epoch_size, t_max=args_opt.epoch_size, eta_min=0.0))
-
- opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr,
- config.momentum, config.weight_decay, loss_scale)
- # 1e-4
- # opt = nn.AdamWeightDecay()
- # optimizer = torch.optim.SGD(model.parameters(), opt.lr, momentum=0.9, nesterov=True)
- params = net.trainable_params()
- # opt = nn.SGD(params=params, momentum=0.9, nesterov=True, learning_rate=1e-4)
- opt = nn.Adam(params=params, momentum=0.9, nesterov=True, learning_rate=1e-4)
- # net = TrainOneStepCell(net, opt, loss_scale) # 将优化器加入网络中
-
- model = Model(net, optimizer=opt)
-
- # cb = [Monitor(lr_init=lr.asnumpy())]
- cb = [LossMonitor()]
- config_ck = CheckpointConfig(save_checkpoint_steps=dataset_size * args_opt.save_checkpoint_epochs,
- keep_checkpoint_max=config.keep_checkpoint_max)
- ckpt_cb = ModelCheckpoint(prefix="EfficientDet", directory=config.save_checkpoint_path, config=config_ck)
- print("============== Starting Training ==============")
- if args_opt.distribute:
- if rank == 0:
- cb += [ckpt_cb]
- # 分析算子的性能
- model.train(args_opt.epoch_size, dataset, callbacks=cb, dataset_sink_mode=False)
-
- else:
- # cb += [ckpt_cb]
- model.train(args_opt.epoch_size, dataset, callbacks=cb, dataset_sink_mode=False)
-
- # profiler.analyse()
-
- print("============== End Training ==============")
-
- if __name__ == '__main__':
- main()
|