|
- """Train retinanet and get checkpoint files."""
-
- import os
- import argparse
- import ast
- import math
- from mindspore import Model, nn, DynamicLossScaleManager
- import mindspore.nn as nn
- 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
- # from utils.utils import init_weights
- from mindspore.common import initializer as initier
- from src.efficientnet.model import EfficientNet
- from src.efficientdet.model import BiFPN,Classifier,Regressor
- from mindspore.common import dtype as mstype
- from mindspore.ops import composite as C
- from mindspore.ops import operations as P
- from mindspore.ops import functional as F
-
- from src.monitor import Monitor
- from mindspore.profiler import Profiler
- # from utils.util_ import init_weights
- # 云上训练代码的改动之处
- import mindspore.common.initializer as weight_init
-
- # import moxing as mox # 将数据拷贝到脚本的包
- set_seed(1)
-
-
- GRADIENT_CLIP_TYPE = 1
- GRADIENT_CLIP_VALUE = 1.0
- clip_grad = C.MultitypeFuncGraph("clip_grad")
- grad_scale = ms.ops.MultitypeFuncGraph("grad_scale")
-
- @grad_scale.register("Tensor", "Tensor")
- def gradient_scale(scale, grad):
- return grad * ms.ops.cast(scale, ms.ops.dtype(grad))
-
-
- @clip_grad.register("Number", "Number", "Tensor")
- def _clip_grad(clip_type, clip_value, grad):
- """
- Clip gradients.
-
- Inputs:
- clip_type (int): The way to clip, 0 for 'value', 1 for 'norm'.
- clip_value (float): Specifies how much to clip.
- grad (tuple[Tensor]): Gradients.
-
- Outputs:
- tuple[Tensor], clipped gradients.
- """
- if clip_type not in (0, 1):
- return grad
- dt = F.dtype(grad)
- if clip_type == 0:
- new_grad = C.clip_by_value(grad, F.cast(F.tuple_to_array((-clip_value,)), dt),
- F.cast(F.tuple_to_array((clip_value,)), dt))
- else:
- new_grad = nn.ClipByNorm()(grad, F.cast(F.tuple_to_array((clip_value,)), dt))
- return new_grad
-
-
- class EfficientDetTrainOneStepCell(nn.TrainOneStepCell):
- """
- Encapsulation class of bert network training.
-
- Append an optimizer to the training network after that the construct
- function can be called to create the backward graph.
-
- Args:
- network (Cell): The training network. Note that loss function should have been added.
- optimizer (Optimizer): Optimizer for updating the weights.
- sens (Number): The adjust parameter. Default: 1.0.
- enable_clip_grad (boolean): If True, clip gradients in BertTrainOneStepCell. Default: True.
- """
-
- def __init__(self, network, optimizer, sens=1.0, enable_clip_grad=True):
- super(EfficientDetTrainOneStepCell, self).__init__(network, optimizer, sens)
- self.cast = P.Cast()
- self.hyper_map = C.HyperMap()
- self.enable_clip_grad = enable_clip_grad
-
- def set_sens(self, value):
- self.sens = value
-
- def construct(self, x, y):
- """Defines the computation performed."""
- weights = self.weights
-
- loss = self.network(x, y)
- grads = self.grad(self.network, weights)(x, y, self.cast(F.tuple_to_array((self.sens,)),
- mstype.float32))
- if self.enable_clip_grad:
- grads = self.hyper_map(F.partial(clip_grad, GRADIENT_CLIP_TYPE, GRADIENT_CLIP_VALUE), grads)
-
- grads = self.grad_reducer(grads)
- succ = self.optimizer(grads)
- return F.depend(loss, succ)
-
- def _calculate_fan_in_and_fan_out(tensor):
- """
- _calculate_fan_in_and_fan_out
- """
- dimensions = len(tensor)
- if dimensions < 2:
- raise ValueError("Fan in and fan out can not be computed for tensor"
- " with fewer than 2 dimensions")
- if dimensions == 2: # Linear
- fan_in = tensor[1]
- fan_out = tensor[0]
- else:
- num_input_fmaps = tensor[1]
- num_output_fmaps = tensor[0]
- receptive_field_size = 1
- if dimensions > 2:
- receptive_field_size = tensor[2] * tensor[3]
- fan_in = num_input_fmaps * receptive_field_size
- fan_out = num_output_fmaps * receptive_field_size
- return fan_in, fan_out
-
-
-
- def init_weights(model):
- # 返回所有模块的迭代器
- for name, cell in model.cells_and_names():
- is_conv_layer = isinstance(cell, nn.Conv2d)
-
- if is_conv_layer:
-
- if "conv_list" in name or "header" in name:
- fan_in, fan_out = _calculate_fan_in_and_fan_out(cell.weight.shape)
- sigma = math.sqrt(1. / float(fan_in)) # 这里计算的是std 而不是bound mu, sigma = 0, 0.1 # 均值和标准差
- data = ms.Tensor(np.random.normal(loc = 0, scale=sigma, size=cell.weight.shape).astype(np.float32))
- cell.weight.set_data(weight_init.initializer(data, cell.weight.shape))
- else:
- cell.weight.set_data(weight_init.initializer(weight_init.HeUniform(),
- cell.weight.shape,
- cell.weight.dtype))
-
- if cell.has_bias is True:
- if "header_cls" in name:
- bias_value = -np.log((1 - 0.01) / 0.01)
- cell.bias.set_data(weight_init.initializer(bias_value, cell.bias.shape))
- else:
- cell.bias.set_data(weight_init.initializer('zeros', cell.bias.shape))
-
-
- class WithLossCell(nn.Cell):
- def __init__(self, backbone, loss):
- super(WithLossCell, self).__init__()
- self.backbone = backbone
- self.loss = loss
-
- def construct(self, x, y):
- _, 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=8, help="Num parallel workers.")
- parser.add_argument("--data_url", type=str, default=None, help="mindrecord dir")
- parser.add_argument("--train_url", type=str, default=None, help="ckpt output dir in obs")
- parser.add_argument("--lr", type=float, default=0.001, 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=2, 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=10, 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.GRAPH_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.")
-
- 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='./profile', is_detail=True, is_show_op_path=False)
- net = EfficientDetBackbone(len(config.coco_classes), compound_coef=0, # 先设置为0
- ratios=eval(config.anchors_ratios), scales=eval(config.anchors_scales)) # 传入参数
-
- init_weights(net)
-
- if args_opt.pre_trained:
- param_dict = load_checkpoint(args_opt.pre_trained)
- load_param_into_net(net, param_dict)
-
- # load_backbone(net, args_opt.pretrained_backbone)
-
- loss = FocalLoss()
- net = WithLossCell(net, loss)
-
- net.set_train()
- # loss_scale = float(args_opt.loss_scale)
-
- loss_scale_manager = DynamicLossScaleManager()
-
- # lr = Tensor(get_lr_cosine(init_lr=0.012, steps_per_epoch=dataset_size, warmup_epochs=int(args_opt.epoch_size / 50),
- # max_epoch=args_opt.epoch_size, t_max=args_opt.epoch_size, eta_min=0.0))
-
- lr = 1e-8
-
- opt = nn.Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr,
- config.momentum, config.weight_decay)
-
- net.set_train()
-
- model = Model(net, loss_scale_manager = loss_scale_manager, optimizer=opt, amp_level="O0")
-
- cb = [LossMonitor(), TimeMonitor()]
-
- 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=True)
-
- else:
- cb += [ckpt_cb]
- model.train(args_opt.epoch_size, dataset, callbacks=cb, dataset_sink_mode=True)
-
- # profiler.analyse()
-
- print("============== End Training ==============")
-
- if __name__ == '__main__':
- main()
|