def step_end(self, run_context):
"""
Save the checkpoint at the end of step.
Args:
run_context (RunContext): Context of the train running.
"""
if _is_role_pserver():
self._prefix = "PServer_" + str(
_get_ps_mode_rank()) + "_" + self._prefix
cb_params = run_context.original_args()
_make_directory(self._directory)
# save graph (only once)
if not self._graph_saved:
graph_file_name = os.path.join(self._directory,
self._prefix + '-graph.meta')
if os.path.isfile(graph_file_name) and context.get_context(
"mode") == context.GRAPH_MODE:
os.remove(graph_file_name)
_save_graph(cb_params.train_network, graph_file_name)
self._graph_saved = True
thread_list = threading.enumerate()
for thread in thread_list:
if thread.getName() == "asyn_save_ckpt":
thread.join()
self._save_ckpt(cb_params)
def step_end(self, run_context):
"""
Save the checkpoint at the end of step.
Args:
run_context (RunContext): Context of the train running.
"""
cb_params = run_context.original_args()
# save graph (only once)
if not self._graph_saved:
graph_file_name = os.path.join(self._directory, self._prefix + '-graph.meta')
_save_graph(cb_params.train_network, graph_file_name)
self._graph_saved = True
self._save_ckpt(cb_params)
def step_end(self, run_context):
"""
Save the checkpoint at the end of step.
Args:
run_context (RunContext): Context of the train running.
"""
if _is_role_pserver():
self._prefix = "PServer_" + str(
_get_ps_mode_rank()) + "_" + self._prefix
cb_params = run_context.original_args()
# save graph (only once)
if not self._graph_saved:
graph_file_name = os.path.join(self._directory,
self._prefix + '-graph.meta')
_save_graph(cb_params.train_network, graph_file_name)
self._graph_saved = True
self._save_ckpt(cb_params)
def test_save_graph():
""" test_exec_save_graph """
class Net1(nn.Cell):
def __init__(self):
super(Net1, self).__init__()
self.add = P.TensorAdd()
def construct(self, x, y):
z = self.add(x, y)
return z
net = Net1()
net.set_train()
out_me_list = []
x = Tensor(np.random.rand(2, 1, 2, 3).astype(np.float32))
y = Tensor(np.array([1.2]).astype(np.float32))
out_put = net(x, y)
_save_graph(network=net, file_name="net-graph.meta")
out_me_list.append(out_put)
def test_save_graph():
""" test_exec_save_graph """
class Net1(nn.Cell):
def __init__(self):
super(Net1, self).__init__()
self.add = P.Add()
def construct(self, x, y):
z = self.add(x, y)
return z
net = Net1()
net.set_train()
out_me_list = []
x = Tensor(np.random.rand(2, 1, 2, 3).astype(np.float32))
y = Tensor(np.array([1.2]).astype(np.float32))
out_put = net(x, y)
output_file = "net-graph.meta"
_save_graph(network=net, file_name=output_file)
out_me_list.append(out_put)
assert os.path.exists(output_file)
os.chmod(output_file, stat.S_IWRITE)
os.remove(output_file)
