def test_all_trains():
ds_train = create_dataset(
os.path.join('/home/workspace/mindspore_dataset/mnist', "train"), 32,
1)
network = LeNet5(10)
net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
net_opt = nn.Momentum(network.trainable_params(), 0.01, 0.9)
time_cb = TimeMonitor(data_size=ds_train.get_dataset_size())
model = Model(network, net_loss, net_opt, metrics={"Accuracy": Accuracy()})
print("============== Starting Training ==============")
model.train(1, ds_train, callbacks=[time_cb, LossMonitor()])
save_checkpoint(self._train_forward_backward,
"gradient_accumulation.ckpt")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description='MindSpore Grad Cumulative Example')
parser.add_argument(
'--device_target',
type=str,
default="GPU",
choices=['GPU'],
help='device where the code will be implemented (default: GPU)')
parser.add_argument('--data_path',
type=str,
default="./Data",
help='path where the dataset is saved')
args = parser.parse_args()
context.set_context(mode=context.GRAPH_MODE,
device_target=args.device_target)
ds_train = create_dataset(os.path.join(args.data_path, "train"), 32)
net = LeNet5(10)
net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction="mean")
net_opt = nn.Momentum(net.trainable_params(), 0.01, 0.9)
model = GradientAccumulation(net, net_loss, net_opt)
print("============== Starting Training ==============")
model.train_process(10, ds_train, mini_steps=4)