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, inputs_, label_) context.reset_auto_parallel_context()
def compile_net_con(net): context.set_context(mode=context.GRAPH_MODE, save_graphs=False) optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) train_net = TrainOneStepCell(net, optimizer) train_net.set_auto_parallel() _executor.compile(train_net, _x_c) context.reset_auto_parallel_context()
def compile_net(net): 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() _executor.compile(train_net, _x, _b) 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, 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 auto_parallel_compile_net(mode, dev_num, 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 = Net2(strategy1, strategy2) net = _VirtualDatasetCell(net) optimizer = Momentum(net.trainable_params(), learning_rate=0.1, momentum=0.9) train_network = TrainOneStepCell(net, optimizer) train_network.set_auto_parallel() _executor.compile(train_network, inputs, label) context.reset_auto_parallel_context()
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 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)