init("nccl")
cfg.rank = get_rank()
cfg.group_size = get_group_size()
parallel_mode = ParallelMode.DATA_PARALLEL
context.set_auto_parallel_context(parallel_mode=parallel_mode, device_num=cfg.group_size,
parameter_broadcast=True, gradients_mean=True)
else:
cfg.rank = 0
cfg.group_size = 1
# dataloader
dataset = create_dataset(args_opt.dataset_path, True, cfg.rank, cfg.group_size)
batches_per_epoch = dataset.get_dataset_size()
print("Batches Per Epoch: ", batches_per_epoch)
# network
net = ShuffleNetV2(n_class=cfg.num_classes, model_size=args_opt.model_size)
# loss
loss = CrossEntropySmooth(sparse=True, reduction="mean",
smooth_factor=cfg.label_smooth_factor, num_classes=cfg.num_classes)
# learning rate schedule
lr = get_lr_basic(lr_init=cfg.lr_init, total_epochs=cfg.epoch_size,
steps_per_epoch=batches_per_epoch, is_stair=True)
lr = Tensor(lr)
# optimizer
decayed_params = []
no_decayed_params = []
for param in net.trainable_params():
if 'beta' not in param.name and 'gamma' not in param.name and 'bias' not in param.name:
def create_network(name, *args, **kwargs):
if name == "shufflenetv2":
return ShuffleNetV2(*args, **kwargs)
raise NotImplementedError(f"{name} is not implemented in the repo")
choices=["AIR", "ONNX", "MINDIR"],
default="AIR",
help="file format")
parser.add_argument(
"--device_target",
type=str,
default="GPU",
choices=["Ascend", "GPU", "CPU"],
help="device where the code will be implemented (default: GPU)")
args = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
if args.device_target == "Ascend":
context.set_context(device_id=args.device_id)
if __name__ == '__main__':
if args.device_target != 'GPU':
raise ValueError("Only supported GPU now.")
net = ShuffleNetV2(n_class=cfg.num_classes)
ckpt = load_checkpoint(args.ckpt_file)
load_param_into_net(net, ckpt)
net.set_train(False)
input_data = Tensor(np.ones([args.batch_size, 3, args.height, args.width]),
ms.float32)
export(net,
input_data,
file_name=args.file_name,
file_format=args.file_format)
def shufflenetv2(*args, **kwargs):
return ShuffleNetV2(*args, **kwargs)