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_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_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 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 auto_parallel_compile_net(mode,
dev_num,
net,
strategy1=None,
strategy2=None):
context.set_context(mode=context.GRAPH_MODE)
context.set_auto_parallel_context(parallel_mode=mode,
device_num=dev_num,
enable_parallel_optimizer=True)
inputs = Tensor(np.ones([32, 48]).astype(np.float32))
label = Tensor(np.zeros([32, 16]).astype(np.float32))
net = net(strategy1, strategy2)
net = _VirtualDatasetCell(net)
optimizer = Momentum(net.trainable_params(),
learning_rate=0.1,
momentum=0.9)
train_network = TrainOneStepCell(net, optimizer).set_comm_fusion(4)
train_network.set_auto_parallel()
train_network.set_train()
_executor.compile(train_network,
inputs,
label,
phase="train",
auto_parallel_mode=True)
context.reset_auto_parallel_context()
return train_network
def compile_net(net):
optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, _x, _b)
context.reset_auto_parallel_context()
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_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_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_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_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 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_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 compile_net(net):
optimizer = Adam(net.trainable_params(),
learning_rate=0.1,
loss_scale=1024.0,
weight_decay=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_train()
_executor.compile(train_net, _x, _b)
def compile_net1(net):
context.set_context(mode=context.GRAPH_MODE, save_graphs=True)
optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, _x1)
context.reset_auto_parallel_context()
def compile_net(net):
context.set_context(save_graphs=False)
optimizer = LazyAdam(net.trainable_params(), learning_rate=0.1)
optimizer.sparse_opt.add_prim_attr("primitive_target", "CPU")
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, _x, _b, auto_parallel_mode=True)
context.reset_auto_parallel_context()
def compile_net(net, x=_b, b=_b):
context.set_context(save_graphs=True)
optimizer = Momentum(net.trainable_params(),
learning_rate=0.1,
momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, x, b)
context.reset_auto_parallel_context()
def compile_net(net, shape):
x = Tensor(np.ones(shape), dtype=ms.int32)
y = Tensor(np.ones(shape), dtype=ms.float32)
z = Tensor(np.ones(shape), dtype=ms.int32)
optimizer = Adam(net.trainable_params(), learning_rate=0.1)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, x, y, z)
context.reset_auto_parallel_context()
def test_async_dump_net_multi_layer_mode1():
test_name = "test_async_dump_net_multi_layer_mode1"
json_file = os.path.join(os.getcwd(), "{}.json".format(test_name))
device_id = context.get_context("device_id")
dump_full_path = os.path.join("/tmp/async_dump/",
"{}_{}".format(test_name, device_id))
os.system("rm -rf {}/*".format(dump_full_path))
os.environ["MINDSPORE_DUMP_CONFIG"] = json_file
weight = Tensor(np.ones((1000, 2048)).astype(np.float32))
bias = Tensor(np.ones((1000, )).astype(np.float32))
net = ReluReduceMeanDenseRelu(weight, bias, 2048, 1000)
criterion = SoftmaxCrossEntropyWithLogits(sparse=False)
optimizer = Momentum(learning_rate=0.1,
momentum=0.1,
params=filter(lambda x: x.requires_grad,
net.get_parameters()))
net_with_criterion = WithLossCell(net, criterion)
train_network = TrainOneStepCell(net_with_criterion, optimizer)
train_network.set_train()
inputs = Tensor(np.random.randn(32, 2048, 7, 7).astype(np.float32))
label = Tensor(np.zeros(shape=(32, 1000)).astype(np.float32))
net_dict = train_network(inputs, label)
dump_path = "/tmp/async_dump/{}/device_{}/test_graph_0/0/0/".format(
test_name, device_id)
dump_file = os.listdir(dump_path)
dump_file_name = ""
for file in dump_file:
if "SoftmaxCrossEntropyWithLogits" in file:
dump_file_name = file
dump_file_full_path = os.path.join(dump_path, dump_file_name)
npy_path = os.path.join(os.getcwd(), "./{}".format(test_name))
if os.path.exists(npy_path):
shutil.rmtree(npy_path)
os.mkdir(npy_path)
tool_path = search_path('/usr/local/Ascend', 'msaccucmp.pyc')
if tool_path:
cmd = "python {0} convert -d {1} -out {2}".format(
tool_path, dump_file_full_path, npy_path)
os.system(cmd)
npy_file_list = os.listdir(npy_path)
dump_result = {}
for file in npy_file_list:
if "output.0.npy" in file:
dump_result["output0"] = np.load(os.path.join(npy_path, file))
for index, value in enumerate(net_dict):
assert value.asnumpy() == dump_result["output0"][index]
else:
print('not find convert tools msaccucmp.pyc')
def train(net, data, label):
learning_rate = 0.01
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()
res = train_network(data, label)
print(res)
assert res
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()
output = train_network(predict, label)
def test_trainTensor(num_classes=10, epoch=15, batch_size=32):
net = AlexNet(num_classes)
lr = 0.1
momentum = 0.9
optimizer = Momentum(filter(lambda x: x.requires_grad, net.get_parameters()), lr, momentum, weight_decay=0.0001)
criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
net_with_criterion = WithLossCell(net, criterion)
train_network = TrainOneStepCell(net_with_criterion, optimizer)
train_network.set_train()
losses = []
for i in range(0, epoch):
data = Tensor(np.ones([batch_size, 3, 227, 227]).astype(np.float32) * 0.01)
label = Tensor(np.ones([batch_size]).astype(np.int32))
loss = train_network(data, label).asnumpy()
losses.append(loss)
assert losses[-1] < 0.01
def test_trainTensor(num_classes=10, epoch=8, batch_size=1):
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
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)
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 train(net, data, label):
learning_rate = 0.01
momentum = 0.9
optimizer = Momentum(
filter(lambda x: x.requires_grad, net.get_parameters()), learning_rate,
momentum)
criterion = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
net_with_criterion = WithLossCell(net, criterion)
train_network = TrainOneStepCell(net_with_criterion,
optimizer) # optimizer
train_network.set_train()
res = train_network(data, label)
print("+++++++++Loss+++++++++++++")
print(res)
print("+++++++++++++++++++++++++++")
diff = res.asnumpy()[0] - 2.3025851
assert np.all(diff < 1.e-6)
def test_train_lenet():
epoch = 100
net = LeNet()
momentum = 0.9
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).asnumpy()
losses.append(loss)
assert losses[-1] < 0.01
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 test_unique_row_split():
class Net(nn.Cell):
def __init__(self):
super().__init__()
self.unique = P.Unique().shard(((1, ), ))
self.relu = P.ReLU()
self.mul = P.Mul()
self.embedding_lookp = P.Gather().shard(((8, 1), (1, )))
self.embedding_table = Parameter(initializer(
'normal', [2000, 128]),
name='embedding_table')
self.gatherv2 = P.Gather().shard(((1, 1), (1, )))
self.reshape = P.Reshape()
self.matmul = P.MatMul()
self.mul_weight = Parameter(Tensor(
np.full([32, 64, 1], 0.5, dtype=np.float32)),
name="mul_weight")
def construct(self, indices):
indices_flatten = self.reshape(indices, (-1, ))
unique_id, unique_idx = self.unique(indices_flatten)
unique_id_weight = self.embedding_lookp(self.embedding_table,
unique_id, 0)
weight_flatten = self.gatherv2(unique_id_weight, unique_idx, 0)
weight = self.reshape(weight_flatten, (32, 64, 128))
vx = self.mul(weight, self.mul_weight)
return vx
size = 8
context.set_auto_parallel_context(device_num=size,
global_rank=0,
parallel_mode="semi_auto_parallel")
x = Tensor(np.ones([32, 64]), dtype=ms.int32)
net = Net()
optimizer = Momentum(net.trainable_params(),
learning_rate=0.1,
momentum=0.9)
train_net = TrainOneStepCell(net, optimizer)
train_net.set_auto_parallel()
train_net.set_train()
_executor.compile(train_net, x)
def test_sit_embedding_lookup_net():
indices = Tensor(np.array([0, 1, 2]).astype(np.int32))
label = Tensor(np.random.randn(3, 8).astype(np.float32))
net1 = NetWithEmbeddingLookUp(vocab_size=8, embedding_size=8, target="CPU")
loss = nn.SoftmaxCrossEntropyWithLogits(reduction="mean")
optimizer1 = nn.Adam(params=net1.trainable_params(), learning_rate=0.1)
optimizer1.unique = True
train_network1 = TrainOneStepCell(WithLossCell(net1, loss), optimizer1)
train_network1.set_train()
out1 = train_network1(indices, label)
net2 = NetWithEmbeddingLookUp(vocab_size=8, embedding_size=8, target="CPU")
optimizer2 = nn.Adam(params=net2.trainable_params(), learning_rate=0.1)
optimizer2.unique = False
optimizer2.target = "CPU"
train_network2 = TrainOneStepCell(WithLossCell(net2, loss), optimizer2)
train_network2.set_train()
out2 = train_network2(indices, label)
assert np.allclose(out1.asnumpy(), out2.asnumpy(), 0.001, 0.001)
def test_parameter_update_float32():
""" test_parameter_update """
net = Net()
loss = nn.SoftmaxCrossEntropyWithLogits()
optimizer = Momentum(net.get_parameters(), 0.01, 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)
# construct and compile update graph
param_lr = train_network.parameters_dict()['learning_rate']
update_network = ParameterUpdate(param_lr)
update_network.phase = 'update_param'
input_lr = Tensor(0.0001, mstype.float32)
_executor.compile(update_network, input_lr)
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.target = 'CPU'
criterion = nn.SoftmaxCrossEntropyWithLogits(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(-5, 15, (32, 3), np.int32))
label = Tensor(np.random.randint(0, 9, (32), np.int32))
if _is_role_pserver():
train_network(data, label)
sys.exit()
else:
loss = train_network(data, label).asnumpy()
losses.append(loss)
print(losses)
return losses
