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