def test_train_with_customized_network(self, *args):
"""Test train with customized network."""
args[0].return_value = 64
train_callback = TrainLineage(SUMMARY_DIR, True, self.user_defined_info)
run_context_customized = self.run_context
del run_context_customized['optimizer']
del run_context_customized['net_outputs']
del run_context_customized['loss_fn']
net = WithLossCell(self.net, self.loss_fn)
net_cap = net
net_cap._cells = {'_backbone': self.net,
'_loss_fn': self.loss_fn}
net = TrainOneStep(net, self.optimizer)
net._cells = {'optimizer': self.optimizer,
'network': net_cap,
'backbone': self.net}
run_context_customized['train_network'] = net
train_callback.begin(RunContext(run_context_customized))
train_callback.end(RunContext(run_context_customized))
LINEAGE_DATA_MANAGER.start_load_data().join()
res = filter_summary_lineage(data_manager=LINEAGE_DATA_MANAGER, search_condition=self._search_condition)
assert res.get('object')[0].get('model_lineage', {}).get('loss_function') \
== 'SoftmaxCrossEntropyWithLogits'
assert res.get('object')[0].get('model_lineage', {}).get('network') == 'ResNet'
assert res.get('object')[0].get('model_lineage', {}).get('optimizer') == 'Momentum'
def bn_common(parallel_mode, train_flag, strategy_loss=None):
context.set_context(mode=context.GRAPH_MODE)
context.set_auto_parallel_context(parallel_mode=parallel_mode,
device_num=8)
learning_rate = 0.1
momentum = 0.9
epoch_size = 2
rank_size = 8
predict = Tensor(np.ones([32, 512]), dtype=ms.float32)
label = Tensor(np.ones([32]), dtype=ms.int32)
dataset = Dataset(predict, label, 2)
net = bn_net()
loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
loss.softmax_cross_entropy.shard(strategy_loss)
opt = Momentum(net.trainable_params(), learning_rate, momentum, 0.0001,
1024 * rank_size)
if not train_flag:
net = WithLossCell(net, loss)
net.set_train()
if parallel_mode == ParallelMode.DATA_PARALLEL:
context.set_auto_parallel_context(parameter_broadcast=True)
model = Model(net, loss, opt)
if train_flag:
model.train(epoch_size, dataset, dataset_sink_mode=False)
else:
model._predict(predict, label)
def test_multi_grads():
context.set_context(mode=context.PYNATIVE_MODE, device_target="Ascend")
sparse = False
inputs_np = np.random.rand(32, 1, 32, 32).astype(np.float32)
labels_np = np.random.randint(10, size=32).astype(np.int32)
inputs_np_2 = np.random.rand(64, 1, 32, 32).astype(np.float32)
labels_np_2 = np.random.randint(10, size=64).astype(np.int32)
if not sparse:
labels_np = np.eye(10)[labels_np].astype(np.float32)
labels_np_2 = np.eye(10)[labels_np_2].astype(np.float32)
net = LeNet()
# grad operation
loss_fn = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=sparse)
with_loss_cell = WithLossCell(net, loss_fn)
grad_all = GradWrapWithLoss(with_loss_cell)
grad_out = grad_all(Tensor(inputs_np), Tensor(labels_np)).asnumpy()
assert np.any(grad_out != 0), 'grad result can not be all zeros'
# train-one-step operation
loss_fn = nn.SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=sparse)
optimizer = Momentum(
filter(lambda x: x.requires_grad, net.get_parameters()), 0.01, 0.9)
loss_net = WithLossCell(net, loss_fn)
train_net = TrainOneStepCell(loss_net, optimizer)
train_net.set_train()
train_net(Tensor(inputs_np_2), Tensor(labels_np_2))
def test_train_with_customized_network(self, *args):
"""Test train with customized network."""
args[0].return_value = 64
train_callback = TrainLineage(SUMMARY_DIR, True,
self.user_defined_info)
run_context_customized = self.run_context
del run_context_customized['optimizer']
del run_context_customized['net_outputs']
del run_context_customized['loss_fn']
net = WithLossCell(self.net, self.loss_fn)
net_cap = net
net_cap._cells = {'_backbone': self.net, '_loss_fn': self.loss_fn}
net = TrainOneStep(net, self.optimizer)
net._cells = {
'optimizer': self.optimizer,
'network': net_cap,
'backbone': self.net
}
run_context_customized['train_network'] = net
train_callback.begin(RunContext(run_context_customized))
train_callback.end(RunContext(run_context_customized))
res = get_summary_lineage(summary_dir=SUMMARY_DIR)
assert res.get('hyper_parameters', {}).get('loss_function') \
== 'SoftmaxCrossEntropyWithLogits'
assert res.get('algorithm', {}).get('network') == 'ResNet'
assert res.get('hyper_parameters', {}).get('optimizer') == 'Momentum'
def test_big_batchSize_with_new_interface(num_classes=10,
epoch=8,
batch_size=338):
net = resnet50(num_classes)
criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
net_with_criterion = WithLossCell(net, criterion)
net_with_criterion.set_train()
weights = ParameterTuple(
filter(lambda x: x.requires_grad, net.get_parameters()))
optimizer = Momentum(weights, 0.1, 0.9)
train_network = ForwardValueAndGrad(network=net_with_criterion,
weights=weights,
get_by_list=True,
sens_param=True,
sens=1.0)
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, grads = train_network(data, label)
grads = F.identity(grads)
optimizer(grads)
losses.append(loss)
assert (losses[-1].asnumpy() < 0.8)
def test_train_lenet_with_new_interface(num_classes=10,
epoch=20,
batch_size=32):
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
network = LeNet5(num_classes)
criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
net_with_criterion = WithLossCell(network, criterion)
net_with_criterion.set_train()
weights = ParameterTuple(network.trainable_params())
optimizer = nn.Momentum(weights, 0.1, 0.9)
train_network = ForwardValueAndGrad(network=net_with_criterion,
weights=weights,
get_by_list=True,
sens_param=True)
losses = []
for i in range(0, epoch):
data = Tensor(
np.ones([batch_size, 1, 32, 32]).astype(np.float32) * 0.01)
label = Tensor(np.ones([batch_size]).astype(np.int32))
sens = Tensor(np.ones([1]).astype(np.float32))
loss, grads = train_network(data, label, sens)
grads = F.identity(grads)
optimizer(grads)
losses.append(loss)
assert losses[-1].asnumpy() < 0.01
assert losses[-1].asnumpy() > 0.001
def test_group_dynamic_2():
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.1, 0.2, 0.3)
default_lr = 0.8
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': conv_params,
'lr': conv_lr
}, {
'params': no_conv_params
}]
net.set_train()
loss = nn.SoftmaxCrossEntropyWithLogits()
opt = RMSProp(group_params, learning_rate=default_lr)
assert opt.is_group is True
assert opt.dynamic_lr is True
for lr, param in zip(opt.learning_rate, opt.parameters):
if 'conv' in param.name:
assert np.all(lr.learning_rate.data.asnumpy() == \
Tensor(np.array(list(conv_lr)).astype(np.float32)).asnumpy())
else:
assert np.all(lr.learning_rate.data.asnumpy() == \
Tensor(np.array([default_lr] * 3).astype(np.float32)).asnumpy())
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepCell(net_with_loss, opt)
_executor.compile(train_network, inputs, label)
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 param in conv_params:
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 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 param in conv_params:
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_lenet_nccl():
context.set_auto_parallel_context(parallel_mode="data_parallel",
gradients_mean=True,
device_num=get_group_size())
net = LeNet()
net.set_train()
learning_rate = multisteplr(epoch, 2)
momentum = 0.9
mom_optimizer = Momentum(
filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate,
momentum)
criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
net_with_criterion = WithLossCell(net, criterion)
train_network = TrainOneStepCell(net_with_criterion, mom_optimizer)
train_network.set_train()
losses = []
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))
start = datetime.datetime.now()
for _ in range(epoch):
for _ in range(mini_batch):
loss = train_network(data, label)
losses.append(loss.asnumpy())
end = datetime.datetime.now()
with open("ms_time.txt", "w") as fo1:
fo1.write("time:")
fo1.write(str(end - start))
with open("ms_loss.txt", "w") as fo2:
fo2.write("loss:")
fo2.write(str(losses[-5:]))
assert losses[-1] < 0.01
def test_adam():
epoch = 3
net = NetAdam()
optimizer = Adam(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()
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
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]
context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
losses2 = []
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)
losses2.append(loss.asnumpy())
assert losses2[0] > losses2[1]
assert losses2[1] > losses2[2]
def test_row_tensor_env_get():
class Loss(nn.Cell):
def __init__(self):
super(Loss, self).__init__()
def construct(self, base, target):
return base
class NetWithSparseGatherV2(nn.Cell):
def __init__(self):
super(NetWithSparseGatherV2, self).__init__()
self.w1 = Parameter(Tensor(np.ones([3, 1, 2]).astype(np.float32)),
name="w1")
self.w2 = Parameter(Tensor(np.ones([2, 1, 2]).astype(np.float32)),
name="w2")
self.gatherv2 = MySparseGatherV2()
self.axis = 0
def construct(self, indices):
return self.gatherv2(self.w1, indices, self.axis) * self.w2
inputs = Tensor(np.array([0, 1]).astype(np.int32))
label = Tensor(np.zeros([2, 1, 2]).astype(np.float32))
net = NetWithSparseGatherV2()
net.set_train()
loss = Loss()
optimizer = AdamWeightDecaySparse(net.trainable_params())
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepCell(net_with_loss, optimizer)
train_network(inputs, label)
def test_SGD():
epoch = 3
net = NetSGD()
learning_rate = 0.1
momentum = 0.9
dampening = 0.0
weight_decay = 0.0
nesterov = True
loss_scale = 1.0
optimizer = SGD(filter(lambda x: x.requires_grad,
net.get_parameters()), learning_rate, momentum,
dampening, weight_decay, nesterov, loss_scale)
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 _ 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)
losses.append(loss.asnumpy())
last_loss = 100.0
for loss in losses:
assert last_loss > loss
last_loss = loss
return losses
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 train_mindspore_impl(self):
input_ = Tensor(np.random.randn(self.batch_num, self.input_channels).astype(np.float32))
weight_np = Tensor(np.random.randn(self.output_channels, self.input_channels).astype(np.float32))
bias = Tensor(np.random.randn(self.output_channels).astype(np.float32))
label_np = np.random.randint(self.output_channels, size=self.batch_num)
label_np_onehot = np.zeros(shape=(self.batch_num, self.output_channels)).astype(np.float32)
label_np_onehot[np.arange(self.batch_num), label_np] = 1.0
label = Tensor(label_np_onehot)
ms_dense = Dense(in_channels=self.input_channels,
out_channels=self.output_channels,
weight_init=weight_np,
bias_init=bias, has_bias=True)
criterion = SoftmaxCrossEntropyWithLogits()
optimizer = nn.Adam(ms_dense.trainable_params(),
learning_rate=1e-3,
beta1=0.9, beta2=0.999, eps=self.epsilon,
use_locking=False,
use_nesterov=False, weight_decay=0.0,
loss_scale=1.0)
net_with_criterion = WithLossCell(ms_dense, criterion)
train_network = TrainOneStepCell(net_with_criterion, optimizer)
train_network.set_train()
print('MS Initialized!')
for _ in range(self.epoch):
train_network(input_, label)
output = ms_dense(input_)
print("===============output=================", output)
return output.asnumpy()
def test_parameter_update_int32_and_tensor():
""" test_parameter_update """
net = Net()
loss = nn.SoftmaxCrossEntropyWithLogits()
optimizer = Momentum(net.get_parameters(), Tensor(np.array([0.1, 0.01, 0.001]), mstype.float32), 0.001)
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepCell(net_with_loss, optimizer)
# compile train graph
train_network.set_train()
inputs = Tensor(np.ones([1, 64]).astype(np.float32))
label = Tensor(np.zeros([1, 10]).astype(np.float32))
_executor.compile(train_network, inputs, label)
# test tensor
param_lr = train_network.parameters_dict()['learning_rate']
update_network = ParameterUpdate(param_lr)
update_network.phase = 'update_param'
input_lr = Tensor(np.array([0.2, 0.02, 0.002]), mstype.float32)
_executor.compile(update_network, input_lr)
# test int32
param_step = train_network.parameters_dict()['global_step']
update_global_step = ParameterUpdate(param_step)
input_step = Tensor(np.array([1000]), mstype.int32)
_executor.compile(update_global_step, input_step)
def do_sparse_embedding(ps=False):
epoch = 10
net = LeNet5(10)
if ps:
net.embedding.embedding_table.set_param_ps()
optimizer = Adam(filter(lambda x: x.requires_grad, net.get_parameters()))
optimizer.sparse_opt.add_prim_attr("primitive_target", "CPU")
criterion = nn.SoftmaxCrossEntropyWithLogits(is_grad=False,
sparse=True,
reduction="mean")
net_with_criterion = WithLossCell(net, criterion)
train_network = TrainOneStepCell(net_with_criterion, optimizer)
train_network.set_train()
losses = []
for _ in range(epoch):
data = Tensor(np.random.randint(0, 15, (32, 3), np.int32))
label = Tensor(np.random.randint(0, 9, (32), np.int32))
if envs.get("MS_ROLE") == "MS_PSERVER":
train_network(data, label)
sys.exit()
else:
loss = train_network(data, label).asnumpy()
losses.append(loss)
print(losses)
return losses
def test_momentum():
epoch = 13
net = MomentumNet()
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()
losses = []
for i 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)
losses.append(loss)
print("================================")
print(losses)
"""
expect output:
[[0.04132498 0.00874167 0.00874167 0.00874167 0.00874167
0.00874167 0.00874167 0.00874167 0.00874167 0.00874167]]
"""
error = np.ones(shape=[1, 10]) * 1.0e-6
return losses
def test_momentum():
epoch = 1
net = MomentumNet()
learning_rate = (0.1, 0.2)
momentum = 0.9
optimizer = Momentum(
filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate,
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 _ 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)
losses.append(loss)
print("================================")
print(losses)
return losses
def test_momentum():
epoch = 3
net = NetMomentum()
learning_rate = initializer(Tensor(np.array([0.01]).astype(np.float32)),
[1])
momentum = initializer(Tensor(np.array([0.9]).astype(np.float32)), [1])
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 _ 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)
losses.append(loss)
_ = np.ones(shape=[1, 10]) * 1.0e-6
return losses
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 = 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_AdamWeightDecaySparse():
""" test_AdamWeightDecaySparse """
context.set_context(mode=context.GRAPH_MODE)
class Loss(nn.Cell):
def __init__(self):
super(Loss, self).__init__()
def construct(self, base, target):
return base
class NetWithSparseGatherV2(nn.Cell):
def __init__(self):
super(NetWithSparseGatherV2, self).__init__()
self.w1 = Parameter(Tensor(np.ones([3, 1, 2]).astype(np.float32)),
name="w1",
sparse_grad="sparse_key_w1")
self.w2 = Parameter(Tensor(np.ones([2, 1, 2]).astype(np.float32)),
name="w2")
self.gatherv2 = P.SparseGatherV2()
self.axis = 0
def construct(self, indices):
return self.gatherv2(self.w1, indices, self.axis) * self.w2
inputs = Tensor(np.array([0, 1]).astype(np.int32))
label = Tensor(np.zeros([2, 1, 2]).astype(np.float32))
net = NetWithSparseGatherV2()
net.set_train()
loss = Loss()
optimizer = AdamWeightDecaySparse(net.trainable_params())
net_with_loss = WithLossCell(net, loss)
train_network = TrainOneStepCell(net_with_loss, optimizer)
_executor.compile(train_network, inputs, label)
def test_checkpoint_save_ckpt_seconds():
"""Test checkpoint save ckpt seconds."""
train_config = CheckpointConfig(save_checkpoint_steps=16,
save_checkpoint_seconds=100,
keep_checkpoint_max=0,
keep_checkpoint_per_n_minutes=1)
ckpt_cb = ModelCheckpoint(config=train_config)
cb_params = _InternalCallbackParam()
net = Net()
loss = nn.SoftmaxCrossEntropyWithLogits()
optim = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
network_ = WithLossCell(net, loss)
_train_network = TrainOneStepCell(network_, optim)
cb_params.train_network = _train_network
cb_params.epoch_num = 10
cb_params.cur_epoch_num = 4
cb_params.cur_step_num = 128
cb_params.batch_num = 32
run_context = RunContext(cb_params)
ckpt_cb.begin(run_context)
ckpt_cb.step_end(run_context)
ckpt_cb2 = ModelCheckpoint(config=train_config)
cb_params.cur_epoch_num = 1
cb_params.cur_step_num = 16
ckpt_cb2.begin(run_context)
ckpt_cb2.step_end(run_context)
def test_save_checkpoint():
"""Test save checkpoint."""
train_config = CheckpointConfig(save_checkpoint_steps=16,
save_checkpoint_seconds=0,
keep_checkpoint_max=5,
keep_checkpoint_per_n_minutes=0)
cb_params = _InternalCallbackParam()
net = Net()
loss = nn.SoftmaxCrossEntropyWithLogits()
optim = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
network_ = WithLossCell(net, loss)
_train_network = TrainOneStepCell(network_, optim)
cb_params.train_network = _train_network
cb_params.epoch_num = 10
cb_params.cur_epoch_num = 5
cb_params.cur_step_num = 0
cb_params.batch_num = 32
ckpoint_cb = ModelCheckpoint(prefix="test_ckpt",
directory='./test_files',
config=train_config)
run_context = RunContext(cb_params)
ckpoint_cb.begin(run_context)
ckpoint_cb.step_end(run_context)
if os.path.exists('./test_files/test_ckpt-model.pkl'):
os.chmod('./test_files/test_ckpt-model.pkl', stat.S_IWRITE)
os.remove('./test_files/test_ckpt-model.pkl')
def test_checkpoint_save_ckpt_with_encryption():
"""Test checkpoint save ckpt with encryption."""
train_config = CheckpointConfig(
save_checkpoint_steps=16,
save_checkpoint_seconds=0,
keep_checkpoint_max=5,
keep_checkpoint_per_n_minutes=0,
enc_key=os.urandom(16),
enc_mode="AES-GCM")
ckpt_cb = ModelCheckpoint(config=train_config)
cb_params = _InternalCallbackParam()
net = Net()
loss = nn.SoftmaxCrossEntropyWithLogits()
optim = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
network_ = WithLossCell(net, loss)
_train_network = TrainOneStepCell(network_, optim)
cb_params.train_network = _train_network
cb_params.epoch_num = 10
cb_params.cur_epoch_num = 5
cb_params.cur_step_num = 160
cb_params.batch_num = 32
run_context = RunContext(cb_params)
ckpt_cb.begin(run_context)
ckpt_cb.step_end(run_context)
ckpt_cb2 = ModelCheckpoint(config=train_config)
cb_params.cur_epoch_num = 1
cb_params.cur_step_num = 15
if platform.system().lower() == "windows":
with pytest.raises(NotImplementedError):
ckpt_cb2.begin(run_context)
ckpt_cb2.step_end(run_context)
else:
ckpt_cb2.begin(run_context)
ckpt_cb2.step_end(run_context)
def test_lenet_pynative_train_net():
""" test_lenet_pynative_train_net """
data = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01)
label = Tensor(np.ones([1, 10]).astype(np.float32))
dout = Tensor(np.ones([1]).astype(np.float32))
iteration_num = 1
verification_step = 0
net = LeNet5()
for i in range(0, iteration_num):
# get the gradients
loss_fn = nn.SoftmaxCrossEntropyWithLogits(is_grad=False)
grad_fn = nn.SoftmaxCrossEntropyWithLogits()
grad_net = WithGradCell(net, grad_fn, sens=dout)
gradients = grad_net(data, label)
# update parameters
opt = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
opt(gradients)
# verification
if i == verification_step:
loss_net = WithLossCell(net, loss_fn)
loss_output = loss_net(data, label)
print("The loss of %s-th iteration is %s" % (i, loss_output.asnumpy()))
def test_loss_wrapper():
""" test_loss_wrapper """
input_data = Tensor(np.ones([1, 64]).astype(np.float32))
input_label = Tensor(np.ones([1, 10]).astype(np.float32))
net = Net()
loss = SoftmaxCrossEntropyWithLogits()
cost = WithLossCell(net, loss)
cost(input_data, input_label)
def test_loss_cell_wrapper():
""" test_loss_cell_wrapper """
data = Tensor(np.ones([1, 1, 32, 32]).astype(np.float32) * 0.01)
label = Tensor(np.ones([1, 10]).astype(np.float32))
net = LeNet5()
loss_fn = nn.SoftmaxCrossEntropyWithLogits()
loss_net = WithLossCell(net, loss_fn)
loss_out = loss_net(data, label)
assert loss_out.asnumpy().dtype == 'float32' or loss_out.asnumpy().dtype == 'float64'
def test_exec_save_checkpoint():
net = Net()
loss = SoftmaxCrossEntropyWithLogits(is_grad=False, sparse=True)
opt = Momentum(net.trainable_params(), 0.0, 0.9, 0.0001, 1024)
loss_net = WithLossCell(net, loss)
train_network = TrainOneStepCell(loss_net, opt)
_exec_save_checkpoint(train_network, ckpoint_file_name="./new_ckpt.ckpt")
load_checkpoint("new_ckpt.ckpt")
def test_save_checkpoint_for_network():
""" test save_checkpoint for network"""
net = Net()
loss = SoftmaxCrossEntropyWithLogits(sparse=True)
opt = Momentum(net.trainable_params(), 0.0, 0.9, 0.0001, 1024)
loss_net = WithLossCell(net, loss)
train_network = TrainOneStepCell(loss_net, opt)
save_checkpoint(train_network, ckpt_file_name="./new_ckpt.ckpt")
load_checkpoint("new_ckpt.ckpt")
