def test_dense_gen_graph(): context.set_context(mode=context.GRAPH_MODE) init() network = Net(512, 128) loss_fn = nn.SoftmaxCrossEntropyWithLogits() optimizer = Momentum(filter(lambda x: x.requires_grad, network.get_parameters()), learning_rate=0.1, momentum=0.9) network = WithLossCell(network, loss_fn) context.reset_auto_parallel_context() context.set_auto_parallel_context(parallel_mode=ParallelMode.HYBRID_PARALLEL, mirror_mean=True, device_num=8) network = TrainOneStepCell(network, optimizer) predict = Tensor(np.ones([64, 512]).astype(np.float32) * 0.01) label = Tensor(np.zeros([64, 32]).astype(np.float32)) network.set_auto_parallel() _executor.compile(network, predict, label)
def test_adam_group1(): """ test_adam_group_lr_and_weight_decay """ inputs = Tensor(np.ones([1, 64]).astype(np.float32)) label = Tensor(np.zeros([1, 10]).astype(np.float32)) net = Net() net.set_train() loss = nn.SoftmaxCrossEntropyWithLogits() net_with_loss = WithLossCell(net, loss) all_params = net.trainable_params() poly_decay_lr = nn.polynomial_decay_lr(0.01, 0.0001, total_step=10, step_per_epoch=1, decay_epoch=3, power=1.0) group_params = [{'params': [all_params[0]], 'lr': poly_decay_lr, 'weight_decay': 0.9}, {'params': [all_params[1]]}] optimizer = nn.Adam(group_params, learning_rate=0.1) train_network = TrainOneStepCell(net_with_loss, optimizer) _executor.compile(train_network, inputs, label)
def test_qat_mobile_train(): net = MobileNetV2(num_class=10) img = Tensor(np.ones((1, 3, 224, 224)).astype(np.float32)) label = Tensor(np.ones((1, 10)).astype(np.float32)) net = qat.convert_quant_network(net, quant_delay=0, bn_fold=False, freeze_bn=10000, weight_bits=8, act_bits=8) loss = nn.SoftmaxCrossEntropyWithLogits(reduction='mean') optimizer = nn.Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) net = nn.WithLossCell(net, loss) net = nn.TrainOneStepCell(net, optimizer) net(img, label)
def test_train_lenet(): epoch = 100 net = LeNet() momentum = initializer(Tensor(np.array([0.9]).astype(np.float32)), [1]) learning_rate = multisteplr(epoch, 30) optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate, momentum) criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True) net_with_criterion = WithLossCell(net, criterion) train_network = TrainOneStepCell(net_with_criterion, optimizer) # optimizer train_network.set_train() losses = [] for i in range(epoch): data = Tensor(np.ones([net.batch_size, 3, 32, 32]).astype(np.float32) * 0.01) label = Tensor(np.ones([net.batch_size]).astype(np.int32)) loss = train_network(data, label) losses.append(loss) print(losses)
def test_hook(): net = LeNet5() optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), 0.1, 0.9) criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=False) net_with_criterion = WithLossCell(net, criterion) train_network = GradWrap(net_with_criterion) train_network.set_train() input_data = Tensor(np.ones([net.batch_size, 1, 32, 32]).astype(np.float32) * 0.01) label = Tensor(np.ones([net.batch_size, net.num_class]).astype(np.float32)) output = net(Tensor(input_data)) loss_output = criterion(output, label) grads = train_network(input_data, label) success = optimizer(grads) assert cell_hook_done assert var_hook_done assert cell_bprop_done print(loss_output.asnumpy())
def test_lars_float_lr(): inputs = Tensor(np.ones([1, 64]).astype(np.float32)) label = Tensor(np.zeros([1, 10]).astype(np.float32)) net = Net() net.set_train() loss = nn.SoftmaxCrossEntropyWithLogits() lr = 0.1 SGD = Momentum(net.trainable_params(), lr, 0.9) optimizer = LARS(SGD, epsilon=1e-08, hyperpara=0.02, decay_filter=lambda x: 'bn' not in x.name, lars_filter=lambda x: 'bn' not in x.name) net_with_loss = WithLossCell(net, loss) train_network = TrainOneStepCell(net_with_loss, optimizer) _executor.compile(train_network, inputs, label)
def train_lenet5(): """ Train lenet5 """ config.ckpt_path = config.output_path context.set_context(mode=context.GRAPH_MODE, device_target=config.device_target) ds_train = create_lenet_dataset(os.path.join(config.data_path, "train"), config.batch_size, num_parallel_workers=1) if ds_train.get_dataset_size() == 0: raise ValueError( "Please check dataset size > 0 and batch_size <= dataset size") network = LeNet5(config.num_classes) net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") net_opt = nn.Momentum(network.trainable_params(), config.lr, config.momentum) time_cb = TimeMonitor(data_size=ds_train.get_dataset_size()) config_ck = CheckpointConfig( save_checkpoint_steps=config.save_checkpoint_steps, keep_checkpoint_max=config.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint( prefix="checkpoint_lenet", directory=None if config.ckpt_path == "" else config.ckpt_path, config=config_ck) if config.device_target != "Ascend": model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) else: model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}, amp_level="O2") print("============== Starting Training ==============") model.train(config.epoch_size, ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()])
def test_membership_inference_eval(): net = Net() loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True) opt = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9) model = Model(network=net, loss_fn=loss, optimizer=opt) inference_model = MembershipInference(model, -1) assert isinstance(inference_model, MembershipInference) batch_size = 16 batches = 1 eval_train = ds.GeneratorDataset(dataset_generator(batch_size, batches), ["image", "label"]) eval_test = ds.GeneratorDataset(dataset_generator(batch_size, batches), ["image", "label"]) metrics = ["precision", "accuracy", "recall"] inference_model.eval(eval_train, eval_test, metrics)
def train(network, net_opt, ds_train, prefix, directory, print_times): net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") model = Model(network, loss_fn=net_loss, optimizer=net_opt, metrics={"acc"}) loss_cb = LossMonitor(per_print_times=print_times) config_ck = CheckpointConfig( save_checkpoint_steps=cfg.save_checkpoint_steps, keep_checkpoint_max=cfg.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint(prefix=prefix, directory=directory, config=config_ck) print("============== Starting Training ==============") model.train(epoch_size, ds_train, callbacks=[ckpoint_cb, loss_cb], dataset_sink_mode=False) return model
def test_trainTensor_amp(num_classes=10, epoch=18, batch_size=16): context.set_context(mode=context.GRAPH_MODE, device_target="GPU", enable_mem_reuse=False, enable_dynamic_memory=False) net = resnet50(num_classes) lr = 0.1 momentum = 0.9 optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr, momentum) criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True) train_network = amp.build_train_network(net, optimizer, criterion, level="O2") train_network.set_train() losses = [] for i in range(0, epoch): data = Tensor(np.ones([batch_size, 3, 224, 224]).astype(np.float32) * 0.01) label = Tensor(np.ones([batch_size]).astype(np.int32)) loss = train_network(data, label) losses.append(loss) assert(losses[-1][0].asnumpy() < 1) assert(losses[-1][1].asnumpy() == False) assert(losses[-1][2].asnumpy() > 1)
def test_membership_inference_object_train(): net = Net() loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True) opt = nn.Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9) model = Model(network=net, loss_fn=loss, optimizer=opt) inference_model = MembershipInference(model, 2) assert isinstance(inference_model, MembershipInference) config = [{ "method": "KNN", "params": { "n_neighbors": [3, 5, 7], } }] ds_train = ds.GeneratorDataset(dataset_generator, ["image", "label"]) ds_test = ds.GeneratorDataset(dataset_generator, ["image", "label"]) inference_model.train(ds_train, ds_test, config)
def test_trainTensor_big_batchSize(num_classes=10, epoch=8, batch_size=338): net = resnet50(num_classes) lr = 0.1 momentum = 0.9 optimizer = Momentum( filter(lambda x: x.requires_grad, net.get_parameters()), lr, momentum) criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') net_with_criterion = WithLossCell(net, criterion) train_network = TrainOneStepCell(net_with_criterion, optimizer) # optimizer train_network.set_train() losses = [] for i in range(0, epoch): data = Tensor( np.ones([batch_size, 3, 224, 224]).astype(np.float32) * 0.01) label = Tensor(np.ones([batch_size]).astype(np.int32)) loss = train_network(data, label) losses.append(loss) assert (losses[-1].asnumpy() < 1)
def test_LSTM(): num_epochs = 5 embed_size = 100 num_hiddens = 100 num_layers = 2 bidirectional = True labels = 2 vocab_size = 252193 max_len = 500 weight = np.ones((vocab_size + 1, embed_size)).astype(np.float32) net = SentimentNet(vocab_size=(vocab_size + 1), embed_size=embed_size, num_hiddens=num_hiddens, num_layers=num_layers, bidirectional=bidirectional, weight=weight, labels=labels, batch_size=batch_size) learning_rate = 0.1 momentum = 0.9 optimizer = Momentum( filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate, momentum) criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True) net_with_criterion = WithLossCell(net, criterion) train_network = TrainOneStepCell(net_with_criterion, optimizer) # optimizer train_network.set_train() train_features = Tensor(np.ones([64, max_len]).astype(np.int32)) train_labels = Tensor(np.ones([ 64, ]).astype(np.int32)[0:64]) losses = [] for epoch in range(num_epochs): loss = train_network(train_features, train_labels) losses.append(loss) print("loss:", loss.asnumpy()) assert (losses[-1].asnumpy() < 0.01)
def test_weight_decay(): inputs = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01) label = Tensor(np.ones([1, 10]).astype(np.float32)) net = LeNet5() conv_weight_decay = 0.8 default_weight_decay = 0.0 conv_params = list( filter(lambda x: 'conv' in x.name, net.trainable_params())) no_conv_params = list( filter(lambda x: 'conv' not in x.name, net.trainable_params())) group_params = [{ 'params': no_conv_params }, { 'params': conv_params, 'weight_decay': conv_weight_decay }, { 'order_params': net.trainable_params() }] net.set_train() loss = nn.SoftmaxCrossEntropyWithLogits() opt = SGD(group_params, learning_rate=0.1, weight_decay=default_weight_decay) assert opt.is_group is True assert opt.is_group_lr is False assert opt.is_group_params_ordered is True for weight_decay, decay_flags, param, order_param in zip( opt.weight_decay, opt.decay_flags, opt.parameters, net.trainable_params()): if 'conv' in param.name: assert weight_decay == conv_weight_decay assert decay_flags is True else: assert weight_decay == default_weight_decay assert decay_flags is False assert param.name == order_param.name net_with_loss = WithLossCell(net, loss) train_network = TrainOneStepCell(net_with_loss, opt) _executor.compile(train_network, inputs, label)
def test_inference_phase(): context.set_auto_parallel_context(device_num=8, global_rank=0) context.set_auto_parallel_context(parallel_mode="auto_parallel") set_cost_model_context(run_phase=1) net = Net(512, 128) predict = Tensor(np.ones([64, 512]).astype(np.float32) * 0.001) label = Tensor(np.ones([64, 128]).astype(np.float32)) loss = nn.SoftmaxCrossEntropyWithLogits() optimizer = Momentum(params=net.trainable_params(), learning_rate=0.1, momentum=0.9) net_with_loss = WithLossCell(net, loss) train_network = TrainOneStepCell(net_with_loss, optimizer) train_network.set_train() train_network.set_auto_parallel() _ = train_network(predict, label)
def train_lenet_quant(): context.set_context(mode=context.GRAPH_MODE, device_target=device_target) cfg = quant_cfg ckpt_path = './ckpt_lenet_noquant-10_1875.ckpt' ds_train = create_dataset(os.path.join(data_path, "train"), cfg.batch_size, 1) step_size = ds_train.get_dataset_size() # define fusion network network = LeNet5Fusion(cfg.num_classes) # load quantization aware network checkpoint param_dict = load_checkpoint(ckpt_path) load_nonquant_param_into_quant_net(network, param_dict) # convert fusion network to quantization aware network quantizer = QuantizationAwareTraining(quant_delay=900, bn_fold=False, per_channel=[True, False], symmetric=[True, False]) network = quantizer.quantize(network) # define network loss net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") # define network optimization net_opt = nn.Momentum(network.trainable_params(), cfg.lr, cfg.momentum) # call back and monitor config_ckpt = CheckpointConfig(save_checkpoint_steps=cfg.epoch_size * step_size, keep_checkpoint_max=cfg.keep_checkpoint_max) ckpt_callback = ModelCheckpoint(prefix="ckpt_lenet_quant", config=config_ckpt) # define model model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) print("============== Starting Training ==============") model.train(cfg['epoch_size'], ds_train, callbacks=[ckpt_callback, LossMonitor()], dataset_sink_mode=True) print("============== End Training ==============")
def test_lenet5_train_step_training_pynative(): """test_lenet5_train_step_training_pynative""" context.set_context(mode=context.PYNATIVE_MODE) context.reset_auto_parallel_context() context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL, device_num=8, mirror_mean=True) predict = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01) label = Tensor(np.zeros([1, 10]).astype(np.float32)) DatasetLenet(predict, label, 2) network = LeNet5() loss_fn = nn.SoftmaxCrossEntropyWithLogits() optimizer = Momentum(network.get_parameters(), learning_rate=0.1, momentum=0.9) Model(network=network, loss_fn=loss_fn, optimizer=optimizer) context.set_context(mode=context.GRAPH_MODE) context.reset_auto_parallel_context() _set_has_initializer(False)
def test_train_and_eval_lenet(): context.set_context(mode=context.GRAPH_MODE, device_target="GPU") network = LeNet5(10) net_loss = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True, reduction="mean") net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9) model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) print("============== Starting Training ==============") ds_train = create_dataset( os.path.join('/home/workspace/mindspore_dataset/mnist', "train"), 32, 1) model.train(1, ds_train, callbacks=[LossMonitor()], dataset_sink_mode=True) print("============== Starting Testing ==============") ds_eval = create_dataset( os.path.join('/home/workspace/mindspore_dataset/mnist', "test"), 32, 1) acc = model.eval(ds_eval, dataset_sink_mode=True) print("============== {} ==============".format(acc))
def train(Net): ds_train, ds_test = create_dataset() # 构建网络 network = Net(cfg.num_classes) # 定义模型的损失函数,优化器 net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") net_opt = nn.Adam(network.trainable_params(), cfg.lr) # 训练模型 model = Model(network, loss_fn=net_loss, optimizer=net_opt, metrics={'acc': Accuracy()}) loss_cb = LossMonitor() print("============== Starting Training ==============") model.train(30, ds_train, callbacks=[loss_cb], dataset_sink_mode=True) # 验证 metric = model.eval(ds_test) print(metric) return model
def test_model_eval_error(): """ test_model_eval_error """ dataset_types = (np.float32, np.float32) dataset_shapes = ((32, 3, 224, 224), (32, 3)) dataset = MindData(size=2, batch_size=32, np_types=dataset_types, output_shapes=dataset_shapes, input_indexs=()) net = nn.ReLU() loss = nn.SoftmaxCrossEntropyWithLogits() context.set_context(mode=context.GRAPH_MODE) model_nometrics = Model(net, loss) with pytest.raises(ValueError): model_nometrics.eval(dataset) model_metrics_empty = Model(net, loss, metrics={}) with pytest.raises(ValueError): model_metrics_empty.eval(dataset)
def test_adam_weight_decay(): epoch = 3 net = NetAdamWeightDecay() optimizer = AdamWeightDecayOp(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate=0.01) criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') net_with_criterion = WithLossCell(net, criterion) train_network = TrainOneStepCell( net_with_criterion, optimizer) train_network.set_train() losses1 = [] for _ in range(epoch): data = Tensor(np.arange(0, 16).reshape( 1, 1, 4, 4).astype(np.float32) * 0.01) label = Tensor(np.array([0]).astype(np.int32)) loss = train_network(data, label) losses1.append(loss.asnumpy()) assert losses1[0] > losses1[1] assert losses1[1] > losses1[2]
def train_wrap(net, loss_fn=None, optimizer=None, weights=None): """train_wrap""" if loss_fn is None: loss_fn = nn.SoftmaxCrossEntropyWithLogits() loss_net = nn.WithLossCell(net, loss_fn) loss_net.set_train() if weights is None: weights = ParameterTuple(net.trainable_params()) if optimizer is None: optimizer = nn.Adam(weights, learning_rate=1e-3, beta1=0.9, beta2=0.999, eps=1e-8, use_locking=False, use_nesterov=False, weight_decay=0.0, loss_scale=1.0) train_net = nn.TrainOneStepCell(loss_net, optimizer) return train_net
def test_trainTensor_amp(num_classes=10, epoch=18, batch_size=16): net = resnet50(num_classes) lr = 0.1 momentum = 0.9 optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr, momentum) criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean') train_network = amp.build_train_network( net, optimizer, criterion, level="O2") train_network.set_train() losses = [] for i in range(0, epoch): data = Tensor(np.ones([batch_size, 3, 224, 224] ).astype(np.float32) * 0.01) label = Tensor(np.ones([batch_size]).astype(np.int32)) loss = train_network(data, label) losses.append(loss) assert (losses[-1][0].asnumpy() < 1) assert not losses[-1][1].asnumpy() assert (losses[-1][2].asnumpy() > 1)
def test_group_lr(): inputs = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01) label = Tensor(np.ones([1, 10]).astype(np.float32)) net = LeNet5() conv_lr = 0.8 default_lr = 0.1 conv_params = list( filter(lambda x: 'conv' in x.name, net.trainable_params())) no_conv_params = list( filter(lambda x: 'conv' not in x.name, net.trainable_params())) group_params = [{ 'params': no_conv_params }, { 'params': conv_params, 'lr': conv_lr }, { 'order_params': net.trainable_params() }] net.set_train() loss = nn.SoftmaxCrossEntropyWithLogits() opt = Momentum(group_params, learning_rate=default_lr, momentum=0.9) assert opt.is_group is True assert opt.is_group_lr is True assert opt.dynamic_lr is False assert opt.is_group_params_ordered is True for lr, param, order_param in zip(opt.learning_rate, opt.parameters, net.trainable_params()): if 'conv' in param.name: assert np.all( lr.data.asnumpy() == Tensor(conv_lr, mstype.float32).asnumpy()) else: assert np.all(lr.data.asnumpy() == Tensor( default_lr, mstype.float32).asnumpy()) assert param.name == order_param.name net_with_loss = WithLossCell(net, loss) train_network = TrainOneStepCell(net_with_loss, opt) _executor.compile(train_network, inputs, label)
def eval_net(): '''eval net''' if config.dataset == 'MR': instance = MovieReview(root_dir=config.data_path, maxlen=config.word_len, split=0.9) elif config.dataset == 'SUBJ': instance = Subjectivity(root_dir=cfg.data_path, maxlen=cfg.word_len, split=0.9) elif config.dataset == 'SST2': instance = SST2(root_dir=config.data_path, maxlen=config.word_len, split=0.9) device_target = config.device_target context.set_context(mode=context.GRAPH_MODE, device_target=config.device_target) if device_target == "Ascend": context.set_context(device_id=get_device_id()) dataset = instance.create_test_dataset(batch_size=config.batch_size) loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True) net = TextCNN(vocab_len=instance.get_dict_len(), word_len=config.word_len, num_classes=config.num_classes, vec_length=config.vec_length) opt = nn.Adam(filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate=0.001, weight_decay=float(config.weight_decay)) param_dict = load_checkpoint(config.checkpoint_file_path) print("load checkpoint from [{}].".format(config.checkpoint_file_path)) load_param_into_net(net, param_dict) net.set_train(False) model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc': Accuracy()}) acc = model.eval(dataset) print("accuracy: ", acc)
def test_adamweightdecay_group(): """ test_adam_group_lr_and_weight_decay """ inputs = Tensor(np.ones([1, 64]).astype(np.float32)) label = Tensor(np.zeros([1, 10]).astype(np.float32)) net = Net() net.set_train() loss = nn.SoftmaxCrossEntropyWithLogits() net_with_loss = WithLossCell(net, loss) all_params = net.trainable_params() schedule_lr = lr_schedules.PolynomialDecayLR(0.01, 0.0001, 3, power=1.0) group_params = [{ 'params': [all_params[0]], 'lr': 0.02, 'weight_decay': 0.9 }, { 'params': [all_params[1]] }] optimizer = nn.AdamWeightDecay(group_params, learning_rate=schedule_lr) train_network = TrainOneStepCell(net_with_loss, optimizer) _executor.compile(train_network, inputs, label)
def test_lamb_group(): """ test_Lamb_group_compile """ inputs = Tensor(np.ones([1, 64]).astype(np.float32)) label = Tensor(np.zeros([1, 10]).astype(np.float32)) net = Net() net.set_train() loss = nn.SoftmaxCrossEntropyWithLogits() warmup_decay_lr = LambLearningRate(0.01, 0.0001, 10, 20, 1.0) all_params = net.trainable_params() group_params = [{ 'params': [all_params[0]], 'lr': warmup_decay_lr, 'weight_decay': 0.9 }, { 'params': [all_params[1]] }] optimizer = Lamb(group_params, 0.02) net_with_loss = WithLossCell(net, loss) train_network = TrainOneStepCell(net_with_loss, optimizer) _executor.compile(train_network, inputs, label)
def __init__(self, config: Dict, model_path: str = None) -> None: super().__init__(config) self.num_labels = self.config.num_labels self.bert = MSBertModel(self.config) self.dropout = nn.Dropout(self.config.hidden_dropout_prob) self.classifier = nn.Dense( self.config.hidden_size, self.num_labels, weight_init=TruncatedNormal(self.config.initializer_range), ) if self.num_labels == 1: self.loss_fct = nn.MSELoss() else: self.loss_fct = nn.SoftmaxCrossEntropyWithLogits( sparse=True, reduction="mean") # self.acc_fct = tf.keras.metrics.SparseCategoricalAccuracy() if model_path is not None: self._load_weights(model_path)
def test_pynative_lenet_train_hook_function_print_and_save_grad(): hook = test_custom_hook_function_base() function = hook.test_custom_hook_function( custom_hook_function_print_and_save_grad, cell_hook_function_print_grad) net = LeNet5(hook_function=function[0], cell_hook_function=function[1]) optimizer = Momentum( filter(lambda x: x.requires_grad, net.get_parameters()), 0.1, 0.9) criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=False) net_with_criterion = WithLossCell(net, criterion) train_network = GradWrap(net_with_criterion) train_network.set_train() input_data = Tensor( np.ones([net.batch_size, 1, 32, 32]).astype(np.float32) * 0.01) label = Tensor(np.ones([net.batch_size, net.num_class]).astype(np.float32)) output = net(Tensor(input_data)) criterion(output, label) grads = train_network(input_data, label) success = optimizer(grads) assert success
def train_lenet(): context.set_context(mode=context.GRAPH_MODE, device_target=device_target) cfg = nonquant_cfg ds_train = create_dataset(os.path.join(data_path, "train"), cfg.batch_size) network = LeNet5(cfg.num_classes) net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean") net_opt = nn.Momentum(network.trainable_params(), cfg.lr, cfg.momentum) time_cb = TimeMonitor(data_size=ds_train.get_dataset_size()) config_ck = CheckpointConfig( save_checkpoint_steps=cfg.save_checkpoint_steps, keep_checkpoint_max=cfg.keep_checkpoint_max) ckpoint_cb = ModelCheckpoint(prefix="ckpt_lenet_noquant", config=config_ck) model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()}) print("============== Starting Training Lenet==============") model.train(cfg['epoch_size'], ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()], dataset_sink_mode=True)