support mlu training

examples/text_classification/rnn/train.py

@@ -27,7 +27,7 @@
 # yapf: disable
 parser = argparse.ArgumentParser(__doc__)
 parser.add_argument("--epochs", type=int, default=15, help="Number of epoches for training.")
-parser.add_argument('--device', choices=['cpu', 'gpu', 'xpu', 'npu'], default="gpu", help="Select which device to train model, defaults to gpu.")
+parser.add_argument('--device', choices=['cpu', 'gpu', 'xpu', 'npu', 'mlu'], default="gpu", help="Select which device to train model, defaults to gpu.")
 parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate used to train.")
 parser.add_argument("--save_dir", type=str, default='checkpoints/', help="Directory to save model checkpoint")
 parser.add_argument("--batch_size", type=int, default=64, help="Total examples' number of a batch for training.")

model_zoo/ernie-1.0/args.py

@@ -92,7 +92,7 @@ def parse_args(MODEL_CLASSES):
     parser.add_argument("--seed", type=int, default=1234, help="Random seed for initialization.")
     parser.add_argument("--num_workers", type=int, default=2, help="Num of workers for DataLoader.")
     parser.add_argument("--check_accuracy", type=str2bool, nargs='?', const=False, help="Check accuracy for training process.")
-    parser.add_argument("--device", type=str, default="gpu", choices=["cpu", "gpu", "xpu"], help="select cpu, gpu, xpu devices.")
+    parser.add_argument("--device", type=str, default="gpu", choices=["cpu", "gpu", "xpu", "mlu"], help="select cpu, gpu, xpu devices.")
     parser.add_argument("--lr_decay_style", type=str, default="cosine", choices=["cosine", "none"], help="Learning rate decay style.")
     parser.add_argument("--share_folder", type=str2bool, nargs='?', const=False, help="Use share folder for data dir and output dir on multi machine.")
 

model_zoo/ernie-1.0/run_pretrain_static.py

@@ -376,7 +376,7 @@ def do_train(args):
                                 args.seed + fleet.worker_index() + 2021)
 
     assert args.device in [
-        "cpu", "gpu", "xpu"
+        "cpu", "gpu", "xpu", "mlu"
     ], "Invalid device! Available device should be cpu, gpu, or xpu."
     place = paddle.set_device(args.device)
 

model_zoo/ernie-3.0/run_qa.py

@@ -132,7 +132,7 @@ def parse_args():
                         help="random seed for initialization")
     parser.add_argument(
         '--device',
-        choices=['cpu', 'gpu', 'xpu'],
+        choices=['cpu', 'gpu', 'xpu', 'mlu'],
         default="gpu",
         help="Select which device to train model, defaults to gpu.")
     parser.add_argument(

model_zoo/ernie-3.0/run_seq_cls.py

@@ -140,6 +140,7 @@ def parse_args():
                         help="random seed for initialization")
     parser.add_argument(
         "--device",
+        choices=['cpu', 'gpu', 'mlu'],
         default="gpu",
         type=str,
         help="The device to select to train the model, is must be cpu/gpu/xpu.")

requirements.txt

@@ -9,4 +9,5 @@ tqdm
 paddlefsl
 sentencepiece
 paddle2onnx
-protobuf>=3.1.0, <=3.20.0
+protobuf>=3.1.0, <=3.20.0
+visualdl