|
- import argparse
- import math
-
- import mindspore.nn as nn
- from mindspore.nn.metrics import Accuracy
- from mindspore import context
- from mindspore.train.callback import LossMonitor, TimeMonitor
- from mindspore.train.model import Model
- from mindspore.nn.loss import SoftmaxCrossEntropyWithLogits
- from mindspore.train.serialization import save_checkpoint
-
- from model import DPCNN
- from dataset import MovieReview
-
- parser = argparse.ArgumentParser(description='Mindspore DPCNN train')
- parser.add_argument('--device_target', type=str, default="CPU", choices=['Ascend', 'GPU', 'CPU'],
- help='device where the code will be implemented (default: Ascend)')
- parser.add_argument('--epoch_size', type=int, default=3, help='train epoches.')
- parser.add_argument('--batch_size', type=int, default=64, help='train batch size.')
- parser.add_argument('--base_lr', type=float, default=0.001, help='init lr')
- parser.add_argument('--data_path', type=str, default="../rt-polaritydata", help='')
- args = parser.parse_args()
-
- if __name__ == '__main__':
- context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
-
- instance = MovieReview(root_dir=args.data_path, maxlen=51, split=0.9)
- train_dataset = instance.create_train_dataset(batch_size=args.batch_size, epoch_size=args.epoch_size)
- batch_num = train_dataset.get_dataset_size()
- test_dataset = instance.create_test_dataset(1)
-
- learning_rate = []
- warm_up = [args.base_lr / math.floor(args.epoch_size / 5) * (i + 1) for _ in range(batch_num) for i in
- range(math.floor(args.epoch_size / 5))]
- shrink = [args.base_lr / (16 * (i + 1)) for _ in range(batch_num) for i in range(math.floor(args.epoch_size * 3 / 5))]
- normal_run = [args.base_lr for _ in range(batch_num) for i in
- range(args.epoch_size - math.floor(args.epoch_size / 5) - math.floor(args.epoch_size * 2 / 5))]
- learning_rate = learning_rate + warm_up + normal_run + shrink
-
- net = DPCNN(instance.get_dict_len())
- opt = nn.Adam(net.trainable_params(), learning_rate=learning_rate, weight_decay=3e-5)
- loss = SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
-
- model = Model(net, loss_fn=loss, optimizer=opt, metrics={'acc': Accuracy()})
-
- time_cb = TimeMonitor(data_size=batch_num)
- loss_cb = LossMonitor()
- print("start train...")
- model.train(args.epoch_size, train_dataset, callbacks=[time_cb, loss_cb], dataset_sink_mode=False)
- print("train success")
-
- print("start eval...")
- train_acc = model.eval(train_dataset, dataset_sink_mode=False)
- test_acc = model.eval(test_dataset, dataset_sink_mode=False)
- save_checkpoint(net, f'DPCNN_valacc_{test_acc["acc"]}.ckpt')
- print(f'train_acc:{train_acc}, test_acc:{test_acc}')
|