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)
