We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[amp] refine transformer fp16 train
### PR types
Performance optimization

### PR changes
Models

### Description
优化transformer动态图fp16速度

1、优化点：

- Adam优化器使用multi_tensor策略

- clear_grad使用set_to_zero=False策略

- dataloader num_workers支持>0 ：--num_workers

- 优化一些代码调用顺序

2、性能测试：

![图片](https://user-images.githubusercontent.com/82555433/148742137-aa3cdeff-333b-4efb-9b9b-5be95c9092dc.png)



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/machine_translation/transformer/train.py`
Content:
```
1 import os
2 import time
3 
4 import yaml
5 import argparse
6 import numpy as np
7 from pprint import pprint
8 from attrdict import AttrDict
9 import inspect
10 
11 import paddle
12 import paddle.distributed as dist
13 
14 import reader
15 from paddlenlp.transformers import TransformerModel, CrossEntropyCriterion
16 from paddlenlp.utils.log import logger
17 
18 from util.record import AverageStatistical
19 from util.to_static import apply_to_static
20 
21 
22 def parse_args():
23     parser = argparse.ArgumentParser()
24     parser.add_argument(
25         "--config",
26         default="./configs/transformer.big.yaml",
27         type=str,
28         help="Path of the config file. ")
29     parser.add_argument(
30         "--benchmark",
31         action="store_true",
32         help="Whether to print logs on each cards and use benchmark vocab. Normally, not necessary to set --benchmark. "
33     )
34     parser.add_argument(
35         "--max_iter",
36         default=None,
37         type=int,
38         help="The maximum iteration for training. ")
39     parser.add_argument(
40         "--train_file",
41         nargs='+',
42         default=None,
43         type=str,
44         help="The files for training, including [source language file, target language file]. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to train. "
45     )
46     parser.add_argument(
47         "--dev_file",
48         nargs='+',
49         default=None,
50         type=str,
51         help="The files for validation, including [source language file, target language file]. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to do validation. "
52     )
53     parser.add_argument(
54         "--vocab_file",
55         default=None,
56         type=str,
57         help="The vocab file. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used."
58     )
59     parser.add_argument(
60         "--unk_token",
61         default=None,
62         type=str,
63         help="The unknown token. It should be provided when use custom vocab_file. "
64     )
65     parser.add_argument(
66         "--bos_token",
67         default=None,
68         type=str,
69         help="The bos token. It should be provided when use custom vocab_file. ")
70     parser.add_argument(
71         "--eos_token",
72         default=None,
73         type=str,
74         help="The eos token. It should be provided when use custom vocab_file. ")
75     args = parser.parse_args()
76     return args
77 
78 
79 def do_train(args):
80     if args.device == "gpu":
81         rank = dist.get_rank()
82         trainer_count = dist.get_world_size()
83     else:
84         rank = 0
85         trainer_count = 1
86         paddle.set_device("cpu")
87 
88     if trainer_count > 1:
89         dist.init_parallel_env()
90 
91     # Set seed for CE
92     random_seed = eval(str(args.random_seed))
93     if random_seed is not None:
94         paddle.seed(random_seed)
95 
96     # Define data loader
97     (train_loader), (eval_loader) = reader.create_data_loader(args)
98 
99     # Define model
100     transformer = TransformerModel(
101         src_vocab_size=args.src_vocab_size,
102         trg_vocab_size=args.trg_vocab_size,
103         max_length=args.max_length + 1,
104         num_encoder_layers=args.n_layer,
105         num_decoder_layers=args.n_layer,
106         n_head=args.n_head,
107         d_model=args.d_model,
108         d_inner_hid=args.d_inner_hid,
109         dropout=args.dropout,
110         weight_sharing=args.weight_sharing,
111         bos_id=args.bos_idx,
112         eos_id=args.eos_idx)
113 
114     transformer = apply_to_static(args, transformer)
115 
116     # Define loss
117     criterion = CrossEntropyCriterion(args.label_smooth_eps, args.bos_idx)
118 
119     scheduler = paddle.optimizer.lr.NoamDecay(
120         args.d_model, args.warmup_steps, args.learning_rate, last_epoch=0)
121 
122     # Define optimizer
123     if 'use_multi_tensor' not in inspect.getfullargspec(
124             paddle.optimizer.Adam.__init__).args:
125         optimizer = paddle.optimizer.Adam(
126             learning_rate=scheduler,
127             beta1=args.beta1,
128             beta2=args.beta2,
129             epsilon=float(args.eps),
130             parameters=transformer.parameters())
131     else:
132         optimizer = paddle.optimizer.Adam(
133             learning_rate=scheduler,
134             beta1=args.beta1,
135             beta2=args.beta2,
136             epsilon=float(args.eps),
137             parameters=transformer.parameters(),
138             use_multi_tensor=True)
139 
140     # Init from some checkpoint, to resume the previous training
141     if args.init_from_checkpoint:
142         model_dict = paddle.load(
143             os.path.join(args.init_from_checkpoint, "transformer.pdparams"))
144         opt_dict = paddle.load(
145             os.path.join(args.init_from_checkpoint, "transformer.pdopt"))
146         transformer.set_state_dict(model_dict)
147         optimizer.set_state_dict(opt_dict)
148         print("loaded from checkpoint.")
149     # Init from some pretrain models, to better solve the current task
150     if args.init_from_pretrain_model:
151         model_dict = paddle.load(
152             os.path.join(args.init_from_pretrain_model, "transformer.pdparams"))
153         transformer.set_state_dict(model_dict)
154         print("loaded from pre-trained model.")
155 
156     # for amp training
157     if args.use_amp:
158         amp_level = 'O2' if args.use_pure_fp16 else 'O1'
159         scaler = paddle.amp.GradScaler(
160             enable=True, init_loss_scaling=args.scale_loss)
161         transformer = paddle.amp.decorate(models=transformer, level=amp_level)
162 
163     # for distributed training
164     if trainer_count > 1:
165         transformer = paddle.DataParallel(transformer)
166 
167     # The best cross-entropy value with label smoothing
168     loss_normalizer = -(
169         (1. - args.label_smooth_eps) * np.log(
170             (1. - args.label_smooth_eps)) + args.label_smooth_eps *
171         np.log(args.label_smooth_eps / (args.trg_vocab_size - 1) + 1e-20))
172 
173     step_idx = 0
174 
175     # For benchmark
176     reader_cost_avg = AverageStatistical()
177     batch_cost_avg = AverageStatistical()
178     batch_ips_avg = AverageStatistical()
179 
180     # Train loop
181     for pass_id in range(args.epoch):
182         epoch_start = time.time()
183 
184         batch_id = 0
185         batch_start = time.time()
186         for input_data in train_loader:
187             train_reader_cost = time.time() - batch_start
188             (src_word, trg_word, lbl_word) = input_data
189 
190             if args.use_amp:
191                 with paddle.amp.auto_cast(
192                         custom_black_list={
193                             'scale', 'reduce_sum', 'elementwise_div'
194                         } if amp_level == 'O2' else {},
195                         level=amp_level):
196                     logits = transformer(src_word=src_word, trg_word=trg_word)
197                     sum_cost, avg_cost, token_num = criterion(logits, lbl_word)
198 
199                 tokens_per_cards = token_num.numpy()
200                 scaled = scaler.scale(avg_cost)  # scale the loss
201                 scaled.backward()  # do backward
202 
203                 scaler.minimize(optimizer, scaled)  # update parameters
204                 if 'set_to_zero' in inspect.getfullargspec(
205                         optimizer.clear_grad).args:
206                     optimizer.clear_grad(set_to_zero=False)
207                 else:
208                     optimizer.clear_grad()
209             else:
210                 logits = transformer(src_word=src_word, trg_word=trg_word)
211                 sum_cost, avg_cost, token_num = criterion(logits, lbl_word)
212                 tokens_per_cards = token_num.numpy()
213 
214                 avg_cost.backward()
215 
216                 optimizer.step()
217                 optimizer.clear_grad()
218 
219             train_batch_cost = time.time() - batch_start
220             reader_cost_avg.record(train_reader_cost)
221             batch_cost_avg.record(train_batch_cost)
222             batch_ips_avg.record(train_batch_cost, tokens_per_cards)
223 
224             # NOTE: For benchmark, loss infomation on all cards will be printed.
225             if step_idx % args.print_step == 0 and (args.benchmark or
226                                                     rank == 0):
227                 total_avg_cost = avg_cost.numpy()
228 
229                 if step_idx == 0:
230                     logger.info(
231                         "step_idx: %d, epoch: %d, batch: %d, avg loss: %f, "
232                         "normalized loss: %f, ppl: %f " %
233                         (step_idx, pass_id, batch_id, total_avg_cost,
234                          total_avg_cost - loss_normalizer,
235                          np.exp([min(total_avg_cost, 100)])))
236                 else:
237                     train_avg_batch_cost = args.print_step / batch_cost_avg.get_total_time(
238                     )
239                     logger.info(
240                         "step_idx: %d, epoch: %d, batch: %d, avg loss: %f, "
241                         "normalized loss: %f, ppl: %f, avg_speed: %.2f step/sec, "
242                         "batch_cost: %.5f sec, reader_cost: %.5f sec, tokens: %d, "
243                         "ips: %.5f words/sec" %
244                         (step_idx, pass_id, batch_id, total_avg_cost,
245                          total_avg_cost - loss_normalizer,
246                          np.exp([min(total_avg_cost, 100)]),
247                          train_avg_batch_cost, batch_cost_avg.get_average(),
248                          reader_cost_avg.get_average(),
249                          batch_ips_avg.get_total_cnt(),
250                          batch_ips_avg.get_average_per_sec()))
251                 reader_cost_avg.reset()
252                 batch_cost_avg.reset()
253                 batch_ips_avg.reset()
254 
255             if step_idx % args.save_step == 0 and step_idx != 0:
256                 # Validation
257                 transformer.eval()
258                 total_sum_cost = 0
259                 total_token_num = 0
260                 with paddle.no_grad():
261                     for input_data in eval_loader:
262                         (src_word, trg_word, lbl_word) = input_data
263                         with paddle.amp.auto_cast(
264                                 custom_black_list={
265                                     'scale', 'reduce_sum', 'elementwise_div'
266                                 } if amp_level == 'O2' else {},
267                                 level=amp_level):
268                             logits = transformer(
269                                 src_word=src_word, trg_word=trg_word)
270                             sum_cost, avg_cost, token_num = criterion(logits,
271                                                                       lbl_word)
272                         total_sum_cost += sum_cost.numpy()
273                         total_token_num += token_num.numpy()
274                         total_avg_cost = total_sum_cost / total_token_num
275                     logger.info("validation, step_idx: %d, avg loss: %f, "
276                                 "normalized loss: %f, ppl: %f" %
277                                 (step_idx, total_avg_cost,
278                                  total_avg_cost - loss_normalizer,
279                                  np.exp([min(total_avg_cost, 100)])))
280                 transformer.train()
281 
282                 if args.save_model and rank == 0:
283                     model_dir = os.path.join(args.save_model,
284                                              "step_" + str(step_idx))
285                     if not os.path.exists(model_dir):
286                         os.makedirs(model_dir)
287                     paddle.save(transformer.state_dict(),
288                                 os.path.join(model_dir, "transformer.pdparams"))
289                     paddle.save(optimizer.state_dict(),
290                                 os.path.join(model_dir, "transformer.pdopt"))
291 
292             #NOTE: Used for benchmark and use None as default. 
293             if args.max_iter and step_idx == args.max_iter:
294                 break
295             batch_id += 1
296             step_idx += 1
297             scheduler.step()
298             batch_start = time.time()
299 
300         #NOTE: Used for benchmark and use None as default. 
301         if args.max_iter and step_idx == args.max_iter:
302             break
303 
304         train_epoch_cost = time.time() - epoch_start
305         logger.info("train epoch: %d, epoch_cost: %.5f s" %
306                     (pass_id, train_epoch_cost))
307 
308     if args.save_model and rank == 0:
309         model_dir = os.path.join(args.save_model, "step_final")
310         if not os.path.exists(model_dir):
311             os.makedirs(model_dir)
312         paddle.save(transformer.state_dict(),
313                     os.path.join(model_dir, "transformer.pdparams"))
314         paddle.save(optimizer.state_dict(),
315                     os.path.join(model_dir, "transformer.pdopt"))
316 
317 
318 if __name__ == "__main__":
319     ARGS = parse_args()
320     yaml_file = ARGS.config
321     with open(yaml_file, 'rt') as f:
322         args = AttrDict(yaml.safe_load(f))
323     args.benchmark = ARGS.benchmark
324     if ARGS.max_iter:
325         args.max_iter = ARGS.max_iter
326     args.train_file = ARGS.train_file
327     args.dev_file = ARGS.dev_file
328     args.vocab_file = ARGS.vocab_file
329     args.unk_token = ARGS.unk_token
330     args.bos_token = ARGS.bos_token
331     args.eos_token = ARGS.eos_token
332     pprint(args)
333 
334     do_train(args)
335
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/machine_translation/transformer/train.py b/examples/machine_translation/transformer/train.py
--- a/examples/machine_translation/transformer/train.py
+++ b/examples/machine_translation/transformer/train.py
@@ -158,7 +158,8 @@
         amp_level = 'O2' if args.use_pure_fp16 else 'O1'
         scaler = paddle.amp.GradScaler(
             enable=True, init_loss_scaling=args.scale_loss)
-        transformer = paddle.amp.decorate(models=transformer, level=amp_level)
+        transformer = paddle.amp.decorate(
+            models=transformer, level=amp_level, save_dtype='float32')
 
     # for distributed training
     if trainer_count > 1:
@@ -260,15 +261,23 @@
                 with paddle.no_grad():
                     for input_data in eval_loader:
                         (src_word, trg_word, lbl_word) = input_data
-                        with paddle.amp.auto_cast(
-                                custom_black_list={
-                                    'scale', 'reduce_sum', 'elementwise_div'
-                                } if amp_level == 'O2' else {},
-                                level=amp_level):
+                        if args.use_amp:
+                            with paddle.amp.auto_cast(
+                                    custom_black_list={
+                                        'scale', 'reduce_sum', 'elementwise_div'
+                                    } if amp_level == 'O2' else {},
+                                    level=amp_level):
+                                logits = transformer(
+                                    src_word=src_word, trg_word=trg_word)
+                                sum_cost, avg_cost, token_num = criterion(
+                                    logits, lbl_word)
+
+                        else:
                             logits = transformer(
                                 src_word=src_word, trg_word=trg_word)
                             sum_cost, avg_cost, token_num = criterion(logits,
                                                                       lbl_word)
+
                         total_sum_cost += sum_cost.numpy()
                         total_token_num += token_num.numpy()
                         total_avg_cost = total_sum_cost / total_token_num

{"golden_diff": "diff --git a/examples/machine_translation/transformer/train.py b/examples/machine_translation/transformer/train.py\n--- a/examples/machine_translation/transformer/train.py\n+++ b/examples/machine_translation/transformer/train.py\n@@ -158,7 +158,8 @@\n         amp_level = 'O2' if args.use_pure_fp16 else 'O1'\n         scaler = paddle.amp.GradScaler(\n             enable=True, init_loss_scaling=args.scale_loss)\n-        transformer = paddle.amp.decorate(models=transformer, level=amp_level)\n+        transformer = paddle.amp.decorate(\n+            models=transformer, level=amp_level, save_dtype='float32')\n \n     # for distributed training\n     if trainer_count > 1:\n@@ -260,15 +261,23 @@\n                 with paddle.no_grad():\n                     for input_data in eval_loader:\n                         (src_word, trg_word, lbl_word) = input_data\n-                        with paddle.amp.auto_cast(\n-                                custom_black_list={\n-                                    'scale', 'reduce_sum', 'elementwise_div'\n-                                } if amp_level == 'O2' else {},\n-                                level=amp_level):\n+                        if args.use_amp:\n+                            with paddle.amp.auto_cast(\n+                                    custom_black_list={\n+                                        'scale', 'reduce_sum', 'elementwise_div'\n+                                    } if amp_level == 'O2' else {},\n+                                    level=amp_level):\n+                                logits = transformer(\n+                                    src_word=src_word, trg_word=trg_word)\n+                                sum_cost, avg_cost, token_num = criterion(\n+                                    logits, lbl_word)\n+\n+                        else:\n                             logits = transformer(\n                                 src_word=src_word, trg_word=trg_word)\n                             sum_cost, avg_cost, token_num = criterion(logits,\n                                                                       lbl_word)\n+\n                         total_sum_cost += sum_cost.numpy()\n                         total_token_num += token_num.numpy()\n                         total_avg_cost = total_sum_cost / total_token_num\n", "issue": "[amp] refine transformer fp16 train\n### PR types\r\nPerformance optimization\r\n\r\n### PR changes\r\nModels\r\n\r\n### Description\r\n\u4f18\u5316transformer\u52a8\u6001\u56fefp16\u901f\u5ea6\r\n\r\n1\u3001\u4f18\u5316\u70b9\uff1a\r\n\r\n- Adam\u4f18\u5316\u5668\u4f7f\u7528multi_tensor\u7b56\u7565\r\n\r\n- clear_grad\u4f7f\u7528set_to_zero=False\u7b56\u7565\r\n\r\n- dataloader num_workers\u652f\u6301>0 \uff1a--num_workers\r\n\r\n- \u4f18\u5316\u4e00\u4e9b\u4ee3\u7801\u8c03\u7528\u987a\u5e8f\r\n\r\n2\u3001\u6027\u80fd\u6d4b\u8bd5\uff1a\r\n\r\n![\u56fe\u7247](https://user-images.githubusercontent.com/82555433/148742137-aa3cdeff-333b-4efb-9b9b-5be95c9092dc.png)\r\n\r\n\n", "before_files": [{"content": "import os\nimport time\n\nimport yaml\nimport argparse\nimport numpy as np\nfrom pprint import pprint\nfrom attrdict import AttrDict\nimport inspect\n\nimport paddle\nimport paddle.distributed as dist\n\nimport reader\nfrom paddlenlp.transformers import TransformerModel, CrossEntropyCriterion\nfrom paddlenlp.utils.log import logger\n\nfrom util.record import AverageStatistical\nfrom util.to_static import apply_to_static\n\n\ndef parse_args():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        \"--config\",\n        default=\"./configs/transformer.big.yaml\",\n        type=str,\n        help=\"Path of the config file. \")\n    parser.add_argument(\n        \"--benchmark\",\n        action=\"store_true\",\n        help=\"Whether to print logs on each cards and use benchmark vocab. Normally, not necessary to set --benchmark. \"\n    )\n    parser.add_argument(\n        \"--max_iter\",\n        default=None,\n        type=int,\n        help=\"The maximum iteration for training. \")\n    parser.add_argument(\n        \"--train_file\",\n        nargs='+',\n        default=None,\n        type=str,\n        help=\"The files for training, including [source language file, target language file]. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to train. \"\n    )\n    parser.add_argument(\n        \"--dev_file\",\n        nargs='+',\n        default=None,\n        type=str,\n        help=\"The files for validation, including [source language file, target language file]. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to do validation. \"\n    )\n    parser.add_argument(\n        \"--vocab_file\",\n        default=None,\n        type=str,\n        help=\"The vocab file. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used.\"\n    )\n    parser.add_argument(\n        \"--unk_token\",\n        default=None,\n        type=str,\n        help=\"The unknown token. It should be provided when use custom vocab_file. \"\n    )\n    parser.add_argument(\n        \"--bos_token\",\n        default=None,\n        type=str,\n        help=\"The bos token. It should be provided when use custom vocab_file. \")\n    parser.add_argument(\n        \"--eos_token\",\n        default=None,\n        type=str,\n        help=\"The eos token. It should be provided when use custom vocab_file. \")\n    args = parser.parse_args()\n    return args\n\n\ndef do_train(args):\n    if args.device == \"gpu\":\n        rank = dist.get_rank()\n        trainer_count = dist.get_world_size()\n    else:\n        rank = 0\n        trainer_count = 1\n        paddle.set_device(\"cpu\")\n\n    if trainer_count > 1:\n        dist.init_parallel_env()\n\n    # Set seed for CE\n    random_seed = eval(str(args.random_seed))\n    if random_seed is not None:\n        paddle.seed(random_seed)\n\n    # Define data loader\n    (train_loader), (eval_loader) = reader.create_data_loader(args)\n\n    # Define model\n    transformer = TransformerModel(\n        src_vocab_size=args.src_vocab_size,\n        trg_vocab_size=args.trg_vocab_size,\n        max_length=args.max_length + 1,\n        num_encoder_layers=args.n_layer,\n        num_decoder_layers=args.n_layer,\n        n_head=args.n_head,\n        d_model=args.d_model,\n        d_inner_hid=args.d_inner_hid,\n        dropout=args.dropout,\n        weight_sharing=args.weight_sharing,\n        bos_id=args.bos_idx,\n        eos_id=args.eos_idx)\n\n    transformer = apply_to_static(args, transformer)\n\n    # Define loss\n    criterion = CrossEntropyCriterion(args.label_smooth_eps, args.bos_idx)\n\n    scheduler = paddle.optimizer.lr.NoamDecay(\n        args.d_model, args.warmup_steps, args.learning_rate, last_epoch=0)\n\n    # Define optimizer\n    if 'use_multi_tensor' not in inspect.getfullargspec(\n            paddle.optimizer.Adam.__init__).args:\n        optimizer = paddle.optimizer.Adam(\n            learning_rate=scheduler,\n            beta1=args.beta1,\n            beta2=args.beta2,\n            epsilon=float(args.eps),\n            parameters=transformer.parameters())\n    else:\n        optimizer = paddle.optimizer.Adam(\n            learning_rate=scheduler,\n            beta1=args.beta1,\n            beta2=args.beta2,\n            epsilon=float(args.eps),\n            parameters=transformer.parameters(),\n            use_multi_tensor=True)\n\n    # Init from some checkpoint, to resume the previous training\n    if args.init_from_checkpoint:\n        model_dict = paddle.load(\n            os.path.join(args.init_from_checkpoint, \"transformer.pdparams\"))\n        opt_dict = paddle.load(\n            os.path.join(args.init_from_checkpoint, \"transformer.pdopt\"))\n        transformer.set_state_dict(model_dict)\n        optimizer.set_state_dict(opt_dict)\n        print(\"loaded from checkpoint.\")\n    # Init from some pretrain models, to better solve the current task\n    if args.init_from_pretrain_model:\n        model_dict = paddle.load(\n            os.path.join(args.init_from_pretrain_model, \"transformer.pdparams\"))\n        transformer.set_state_dict(model_dict)\n        print(\"loaded from pre-trained model.\")\n\n    # for amp training\n    if args.use_amp:\n        amp_level = 'O2' if args.use_pure_fp16 else 'O1'\n        scaler = paddle.amp.GradScaler(\n            enable=True, init_loss_scaling=args.scale_loss)\n        transformer = paddle.amp.decorate(models=transformer, level=amp_level)\n\n    # for distributed training\n    if trainer_count > 1:\n        transformer = paddle.DataParallel(transformer)\n\n    # The best cross-entropy value with label smoothing\n    loss_normalizer = -(\n        (1. - args.label_smooth_eps) * np.log(\n            (1. - args.label_smooth_eps)) + args.label_smooth_eps *\n        np.log(args.label_smooth_eps / (args.trg_vocab_size - 1) + 1e-20))\n\n    step_idx = 0\n\n    # For benchmark\n    reader_cost_avg = AverageStatistical()\n    batch_cost_avg = AverageStatistical()\n    batch_ips_avg = AverageStatistical()\n\n    # Train loop\n    for pass_id in range(args.epoch):\n        epoch_start = time.time()\n\n        batch_id = 0\n        batch_start = time.time()\n        for input_data in train_loader:\n            train_reader_cost = time.time() - batch_start\n            (src_word, trg_word, lbl_word) = input_data\n\n            if args.use_amp:\n                with paddle.amp.auto_cast(\n                        custom_black_list={\n                            'scale', 'reduce_sum', 'elementwise_div'\n                        } if amp_level == 'O2' else {},\n                        level=amp_level):\n                    logits = transformer(src_word=src_word, trg_word=trg_word)\n                    sum_cost, avg_cost, token_num = criterion(logits, lbl_word)\n\n                tokens_per_cards = token_num.numpy()\n                scaled = scaler.scale(avg_cost)  # scale the loss\n                scaled.backward()  # do backward\n\n                scaler.minimize(optimizer, scaled)  # update parameters\n                if 'set_to_zero' in inspect.getfullargspec(\n                        optimizer.clear_grad).args:\n                    optimizer.clear_grad(set_to_zero=False)\n                else:\n                    optimizer.clear_grad()\n            else:\n                logits = transformer(src_word=src_word, trg_word=trg_word)\n                sum_cost, avg_cost, token_num = criterion(logits, lbl_word)\n                tokens_per_cards = token_num.numpy()\n\n                avg_cost.backward()\n\n                optimizer.step()\n                optimizer.clear_grad()\n\n            train_batch_cost = time.time() - batch_start\n            reader_cost_avg.record(train_reader_cost)\n            batch_cost_avg.record(train_batch_cost)\n            batch_ips_avg.record(train_batch_cost, tokens_per_cards)\n\n            # NOTE: For benchmark, loss infomation on all cards will be printed.\n            if step_idx % args.print_step == 0 and (args.benchmark or\n                                                    rank == 0):\n                total_avg_cost = avg_cost.numpy()\n\n                if step_idx == 0:\n                    logger.info(\n                        \"step_idx: %d, epoch: %d, batch: %d, avg loss: %f, \"\n                        \"normalized loss: %f, ppl: %f \" %\n                        (step_idx, pass_id, batch_id, total_avg_cost,\n                         total_avg_cost - loss_normalizer,\n                         np.exp([min(total_avg_cost, 100)])))\n                else:\n                    train_avg_batch_cost = args.print_step / batch_cost_avg.get_total_time(\n                    )\n                    logger.info(\n                        \"step_idx: %d, epoch: %d, batch: %d, avg loss: %f, \"\n                        \"normalized loss: %f, ppl: %f, avg_speed: %.2f step/sec, \"\n                        \"batch_cost: %.5f sec, reader_cost: %.5f sec, tokens: %d, \"\n                        \"ips: %.5f words/sec\" %\n                        (step_idx, pass_id, batch_id, total_avg_cost,\n                         total_avg_cost - loss_normalizer,\n                         np.exp([min(total_avg_cost, 100)]),\n                         train_avg_batch_cost, batch_cost_avg.get_average(),\n                         reader_cost_avg.get_average(),\n                         batch_ips_avg.get_total_cnt(),\n                         batch_ips_avg.get_average_per_sec()))\n                reader_cost_avg.reset()\n                batch_cost_avg.reset()\n                batch_ips_avg.reset()\n\n            if step_idx % args.save_step == 0 and step_idx != 0:\n                # Validation\n                transformer.eval()\n                total_sum_cost = 0\n                total_token_num = 0\n                with paddle.no_grad():\n                    for input_data in eval_loader:\n                        (src_word, trg_word, lbl_word) = input_data\n                        with paddle.amp.auto_cast(\n                                custom_black_list={\n                                    'scale', 'reduce_sum', 'elementwise_div'\n                                } if amp_level == 'O2' else {},\n                                level=amp_level):\n                            logits = transformer(\n                                src_word=src_word, trg_word=trg_word)\n                            sum_cost, avg_cost, token_num = criterion(logits,\n                                                                      lbl_word)\n                        total_sum_cost += sum_cost.numpy()\n                        total_token_num += token_num.numpy()\n                        total_avg_cost = total_sum_cost / total_token_num\n                    logger.info(\"validation, step_idx: %d, avg loss: %f, \"\n                                \"normalized loss: %f, ppl: %f\" %\n                                (step_idx, total_avg_cost,\n                                 total_avg_cost - loss_normalizer,\n                                 np.exp([min(total_avg_cost, 100)])))\n                transformer.train()\n\n                if args.save_model and rank == 0:\n                    model_dir = os.path.join(args.save_model,\n                                             \"step_\" + str(step_idx))\n                    if not os.path.exists(model_dir):\n                        os.makedirs(model_dir)\n                    paddle.save(transformer.state_dict(),\n                                os.path.join(model_dir, \"transformer.pdparams\"))\n                    paddle.save(optimizer.state_dict(),\n                                os.path.join(model_dir, \"transformer.pdopt\"))\n\n            #NOTE: Used for benchmark and use None as default. \n            if args.max_iter and step_idx == args.max_iter:\n                break\n            batch_id += 1\n            step_idx += 1\n            scheduler.step()\n            batch_start = time.time()\n\n        #NOTE: Used for benchmark and use None as default. \n        if args.max_iter and step_idx == args.max_iter:\n            break\n\n        train_epoch_cost = time.time() - epoch_start\n        logger.info(\"train epoch: %d, epoch_cost: %.5f s\" %\n                    (pass_id, train_epoch_cost))\n\n    if args.save_model and rank == 0:\n        model_dir = os.path.join(args.save_model, \"step_final\")\n        if not os.path.exists(model_dir):\n            os.makedirs(model_dir)\n        paddle.save(transformer.state_dict(),\n                    os.path.join(model_dir, \"transformer.pdparams\"))\n        paddle.save(optimizer.state_dict(),\n                    os.path.join(model_dir, \"transformer.pdopt\"))\n\n\nif __name__ == \"__main__\":\n    ARGS = parse_args()\n    yaml_file = ARGS.config\n    with open(yaml_file, 'rt') as f:\n        args = AttrDict(yaml.safe_load(f))\n    args.benchmark = ARGS.benchmark\n    if ARGS.max_iter:\n        args.max_iter = ARGS.max_iter\n    args.train_file = ARGS.train_file\n    args.dev_file = ARGS.dev_file\n    args.vocab_file = ARGS.vocab_file\n    args.unk_token = ARGS.unk_token\n    args.bos_token = ARGS.bos_token\n    args.eos_token = ARGS.eos_token\n    pprint(args)\n\n    do_train(args)\n", "path": "examples/machine_translation/transformer/train.py"}], "after_files": [{"content": "import os\nimport time\n\nimport yaml\nimport argparse\nimport numpy as np\nfrom pprint import pprint\nfrom attrdict import AttrDict\nimport inspect\n\nimport paddle\nimport paddle.distributed as dist\n\nimport reader\nfrom paddlenlp.transformers import TransformerModel, CrossEntropyCriterion\nfrom paddlenlp.utils.log import logger\n\nfrom util.record import AverageStatistical\nfrom util.to_static import apply_to_static\n\n\ndef parse_args():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        \"--config\",\n        default=\"./configs/transformer.big.yaml\",\n        type=str,\n        help=\"Path of the config file. \")\n    parser.add_argument(\n        \"--benchmark\",\n        action=\"store_true\",\n        help=\"Whether to print logs on each cards and use benchmark vocab. Normally, not necessary to set --benchmark. \"\n    )\n    parser.add_argument(\n        \"--max_iter\",\n        default=None,\n        type=int,\n        help=\"The maximum iteration for training. \")\n    parser.add_argument(\n        \"--train_file\",\n        nargs='+',\n        default=None,\n        type=str,\n        help=\"The files for training, including [source language file, target language file]. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to train. \"\n    )\n    parser.add_argument(\n        \"--dev_file\",\n        nargs='+',\n        default=None,\n        type=str,\n        help=\"The files for validation, including [source language file, target language file]. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used to do validation. \"\n    )\n    parser.add_argument(\n        \"--vocab_file\",\n        default=None,\n        type=str,\n        help=\"The vocab file. Normally, it shouldn't be set and in this case, the default WMT14 dataset will be used.\"\n    )\n    parser.add_argument(\n        \"--unk_token\",\n        default=None,\n        type=str,\n        help=\"The unknown token. It should be provided when use custom vocab_file. \"\n    )\n    parser.add_argument(\n        \"--bos_token\",\n        default=None,\n        type=str,\n        help=\"The bos token. It should be provided when use custom vocab_file. \")\n    parser.add_argument(\n        \"--eos_token\",\n        default=None,\n        type=str,\n        help=\"The eos token. It should be provided when use custom vocab_file. \")\n    args = parser.parse_args()\n    return args\n\n\ndef do_train(args):\n    if args.device == \"gpu\":\n        rank = dist.get_rank()\n        trainer_count = dist.get_world_size()\n    else:\n        rank = 0\n        trainer_count = 1\n        paddle.set_device(\"cpu\")\n\n    if trainer_count > 1:\n        dist.init_parallel_env()\n\n    # Set seed for CE\n    random_seed = eval(str(args.random_seed))\n    if random_seed is not None:\n        paddle.seed(random_seed)\n\n    # Define data loader\n    (train_loader), (eval_loader) = reader.create_data_loader(args)\n\n    # Define model\n    transformer = TransformerModel(\n        src_vocab_size=args.src_vocab_size,\n        trg_vocab_size=args.trg_vocab_size,\n        max_length=args.max_length + 1,\n        num_encoder_layers=args.n_layer,\n        num_decoder_layers=args.n_layer,\n        n_head=args.n_head,\n        d_model=args.d_model,\n        d_inner_hid=args.d_inner_hid,\n        dropout=args.dropout,\n        weight_sharing=args.weight_sharing,\n        bos_id=args.bos_idx,\n        eos_id=args.eos_idx)\n\n    transformer = apply_to_static(args, transformer)\n\n    # Define loss\n    criterion = CrossEntropyCriterion(args.label_smooth_eps, args.bos_idx)\n\n    scheduler = paddle.optimizer.lr.NoamDecay(\n        args.d_model, args.warmup_steps, args.learning_rate, last_epoch=0)\n\n    # Define optimizer\n    if 'use_multi_tensor' not in inspect.getfullargspec(\n            paddle.optimizer.Adam.__init__).args:\n        optimizer = paddle.optimizer.Adam(\n            learning_rate=scheduler,\n            beta1=args.beta1,\n            beta2=args.beta2,\n            epsilon=float(args.eps),\n            parameters=transformer.parameters())\n    else:\n        optimizer = paddle.optimizer.Adam(\n            learning_rate=scheduler,\n            beta1=args.beta1,\n            beta2=args.beta2,\n            epsilon=float(args.eps),\n            parameters=transformer.parameters(),\n            use_multi_tensor=True)\n\n    # Init from some checkpoint, to resume the previous training\n    if args.init_from_checkpoint:\n        model_dict = paddle.load(\n            os.path.join(args.init_from_checkpoint, \"transformer.pdparams\"))\n        opt_dict = paddle.load(\n            os.path.join(args.init_from_checkpoint, \"transformer.pdopt\"))\n        transformer.set_state_dict(model_dict)\n        optimizer.set_state_dict(opt_dict)\n        print(\"loaded from checkpoint.\")\n    # Init from some pretrain models, to better solve the current task\n    if args.init_from_pretrain_model:\n        model_dict = paddle.load(\n            os.path.join(args.init_from_pretrain_model, \"transformer.pdparams\"))\n        transformer.set_state_dict(model_dict)\n        print(\"loaded from pre-trained model.\")\n\n    # for amp training\n    if args.use_amp:\n        amp_level = 'O2' if args.use_pure_fp16 else 'O1'\n        scaler = paddle.amp.GradScaler(\n            enable=True, init_loss_scaling=args.scale_loss)\n        transformer = paddle.amp.decorate(\n            models=transformer, level=amp_level, save_dtype='float32')\n\n    # for distributed training\n    if trainer_count > 1:\n        transformer = paddle.DataParallel(transformer)\n\n    # The best cross-entropy value with label smoothing\n    loss_normalizer = -(\n        (1. - args.label_smooth_eps) * np.log(\n            (1. - args.label_smooth_eps)) + args.label_smooth_eps *\n        np.log(args.label_smooth_eps / (args.trg_vocab_size - 1) + 1e-20))\n\n    step_idx = 0\n\n    # For benchmark\n    reader_cost_avg = AverageStatistical()\n    batch_cost_avg = AverageStatistical()\n    batch_ips_avg = AverageStatistical()\n\n    # Train loop\n    for pass_id in range(args.epoch):\n        epoch_start = time.time()\n\n        batch_id = 0\n        batch_start = time.time()\n        for input_data in train_loader:\n            train_reader_cost = time.time() - batch_start\n            (src_word, trg_word, lbl_word) = input_data\n\n            if args.use_amp:\n                with paddle.amp.auto_cast(\n                        custom_black_list={\n                            'scale', 'reduce_sum', 'elementwise_div'\n                        } if amp_level == 'O2' else {},\n                        level=amp_level):\n                    logits = transformer(src_word=src_word, trg_word=trg_word)\n                    sum_cost, avg_cost, token_num = criterion(logits, lbl_word)\n\n                tokens_per_cards = token_num.numpy()\n                scaled = scaler.scale(avg_cost)  # scale the loss\n                scaled.backward()  # do backward\n\n                scaler.minimize(optimizer, scaled)  # update parameters\n                if 'set_to_zero' in inspect.getfullargspec(\n                        optimizer.clear_grad).args:\n                    optimizer.clear_grad(set_to_zero=False)\n                else:\n                    optimizer.clear_grad()\n            else:\n                logits = transformer(src_word=src_word, trg_word=trg_word)\n                sum_cost, avg_cost, token_num = criterion(logits, lbl_word)\n                tokens_per_cards = token_num.numpy()\n\n                avg_cost.backward()\n\n                optimizer.step()\n                optimizer.clear_grad()\n\n            train_batch_cost = time.time() - batch_start\n            reader_cost_avg.record(train_reader_cost)\n            batch_cost_avg.record(train_batch_cost)\n            batch_ips_avg.record(train_batch_cost, tokens_per_cards)\n\n            # NOTE: For benchmark, loss infomation on all cards will be printed.\n            if step_idx % args.print_step == 0 and (args.benchmark or\n                                                    rank == 0):\n                total_avg_cost = avg_cost.numpy()\n\n                if step_idx == 0:\n                    logger.info(\n                        \"step_idx: %d, epoch: %d, batch: %d, avg loss: %f, \"\n                        \"normalized loss: %f, ppl: %f \" %\n                        (step_idx, pass_id, batch_id, total_avg_cost,\n                         total_avg_cost - loss_normalizer,\n                         np.exp([min(total_avg_cost, 100)])))\n                else:\n                    train_avg_batch_cost = args.print_step / batch_cost_avg.get_total_time(\n                    )\n                    logger.info(\n                        \"step_idx: %d, epoch: %d, batch: %d, avg loss: %f, \"\n                        \"normalized loss: %f, ppl: %f, avg_speed: %.2f step/sec, \"\n                        \"batch_cost: %.5f sec, reader_cost: %.5f sec, tokens: %d, \"\n                        \"ips: %.5f words/sec\" %\n                        (step_idx, pass_id, batch_id, total_avg_cost,\n                         total_avg_cost - loss_normalizer,\n                         np.exp([min(total_avg_cost, 100)]),\n                         train_avg_batch_cost, batch_cost_avg.get_average(),\n                         reader_cost_avg.get_average(),\n                         batch_ips_avg.get_total_cnt(),\n                         batch_ips_avg.get_average_per_sec()))\n                reader_cost_avg.reset()\n                batch_cost_avg.reset()\n                batch_ips_avg.reset()\n\n            if step_idx % args.save_step == 0 and step_idx != 0:\n                # Validation\n                transformer.eval()\n                total_sum_cost = 0\n                total_token_num = 0\n                with paddle.no_grad():\n                    for input_data in eval_loader:\n                        (src_word, trg_word, lbl_word) = input_data\n                        if args.use_amp:\n                            with paddle.amp.auto_cast(\n                                    custom_black_list={\n                                        'scale', 'reduce_sum', 'elementwise_div'\n                                    } if amp_level == 'O2' else {},\n                                    level=amp_level):\n                                logits = transformer(\n                                    src_word=src_word, trg_word=trg_word)\n                                sum_cost, avg_cost, token_num = criterion(\n                                    logits, lbl_word)\n\n                        else:\n                            logits = transformer(\n                                src_word=src_word, trg_word=trg_word)\n                            sum_cost, avg_cost, token_num = criterion(logits,\n                                                                      lbl_word)\n\n                        total_sum_cost += sum_cost.numpy()\n                        total_token_num += token_num.numpy()\n                        total_avg_cost = total_sum_cost / total_token_num\n                    logger.info(\"validation, step_idx: %d, avg loss: %f, \"\n                                \"normalized loss: %f, ppl: %f\" %\n                                (step_idx, total_avg_cost,\n                                 total_avg_cost - loss_normalizer,\n                                 np.exp([min(total_avg_cost, 100)])))\n                transformer.train()\n\n                if args.save_model and rank == 0:\n                    model_dir = os.path.join(args.save_model,\n                                             \"step_\" + str(step_idx))\n                    if not os.path.exists(model_dir):\n                        os.makedirs(model_dir)\n                    paddle.save(transformer.state_dict(),\n                                os.path.join(model_dir, \"transformer.pdparams\"))\n                    paddle.save(optimizer.state_dict(),\n                                os.path.join(model_dir, \"transformer.pdopt\"))\n\n            #NOTE: Used for benchmark and use None as default. \n            if args.max_iter and step_idx == args.max_iter:\n                break\n            batch_id += 1\n            step_idx += 1\n            scheduler.step()\n            batch_start = time.time()\n\n        #NOTE: Used for benchmark and use None as default. \n        if args.max_iter and step_idx == args.max_iter:\n            break\n\n        train_epoch_cost = time.time() - epoch_start\n        logger.info(\"train epoch: %d, epoch_cost: %.5f s\" %\n                    (pass_id, train_epoch_cost))\n\n    if args.save_model and rank == 0:\n        model_dir = os.path.join(args.save_model, \"step_final\")\n        if not os.path.exists(model_dir):\n            os.makedirs(model_dir)\n        paddle.save(transformer.state_dict(),\n                    os.path.join(model_dir, \"transformer.pdparams\"))\n        paddle.save(optimizer.state_dict(),\n                    os.path.join(model_dir, \"transformer.pdopt\"))\n\n\nif __name__ == \"__main__\":\n    ARGS = parse_args()\n    yaml_file = ARGS.config\n    with open(yaml_file, 'rt') as f:\n        args = AttrDict(yaml.safe_load(f))\n    args.benchmark = ARGS.benchmark\n    if ARGS.max_iter:\n        args.max_iter = ARGS.max_iter\n    args.train_file = ARGS.train_file\n    args.dev_file = ARGS.dev_file\n    args.vocab_file = ARGS.vocab_file\n    args.unk_token = ARGS.unk_token\n    args.bos_token = ARGS.bos_token\n    args.eos_token = ARGS.eos_token\n    pprint(args)\n\n    do_train(args)\n", "path": "examples/machine_translation/transformer/train.py"}]}