Add AdamW optimizer support for World Language Model example (#1380)

Author: likejazz

word_language_model/README.md

@@ -8,9 +8,11 @@ python main.py --accel --epochs 6 # Train a LSTM on Wikitext-2.
 python main.py --accel --epochs 6 --tied # Train a tied LSTM on Wikitext-2.
 python main.py --accel --tied # Train a tied LSTM on Wikitext-2for 40 epochs.
 python main.py --accel --epochs 6 --model Transformer --lr 5
- # Train a Transformer model on Wikitext-2.
+ # Train a Transformer model on Wikitext-2.
+python main.py --accel --epochs 6 --model Transformer --use-optimizer --lr 0.001
+ # Train a Transformer model with AdamW optimizer on Wikitext-2.
 
-python generate.py --accel  # Generate samples from the default model checkpoint.
+python generate.py --accel # Generate samples from the default model checkpoint.
 ```
 
 > [!NOTE] 
@@ -45,6 +47,7 @@ optional arguments:
  path to export the final model in onnx format
  --nhead NHEAD the number of heads in the encoder/decoder of the transformer model
  --dry-run verify the code and the model
+ --use-optimizer specify whether to use an AdamW optimizer
 ```
 
 With these arguments, a variety of models can be tested.

word_language_model/main.py

@@ -47,7 +47,10 @@
  help='the number of heads in the encoder/decoder of the transformer model')
 parser.add_argument('--dry-run', action='store_true',
  help='verify the code and the model')
-parser.add_argument('--accel', action='store_true',help='Enables accelerated training')
+parser.add_argument('--accel', action='store_true',
+ help='Enables accelerated training')
+parser.add_argument('--use-optimizer', action='store_true',
+ help='Uses AdamW optimizer for gradient updating')
 args = parser.parse_args()
 
 # Set the random seed manually for reproducibility.
@@ -104,6 +107,8 @@ def batchify(data, bsz):
  model = RNNModel(args.model, ntokens, args.emsize, args.nhid, args.nlayers, args.dropout, args.tied).to(device)
 
 criterion = nn.NLLLoss()
+if args.use_optimizer:
+ optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
 
 ###############################################################################
 # Training code
@@ -167,7 +172,10 @@ def train():
  data, targets = get_batch(train_data, i)
  # Starting each batch, we detach the hidden state from how it was previously produced.
  # If we didn't, the model would try backpropagating all the way to start of the dataset.
- model.zero_grad()
+ if args.use_optimizer:
+ optimizer.zero_grad()
+ else:
+ model.zero_grad()
  if args.model == 'Transformer':
  output = model(data)
  output = output.view(-1, ntokens)
@@ -179,8 +187,11 @@ def train():
 
  # `clip_grad_norm` helps prevent the exploding gradient problem in RNNs / LSTMs.
  torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip)
- for p in model.parameters():
- p.data.add_(p.grad, alpha=-lr)
+ if args.use_optimizer:
+ optimizer.step()
+ else:
+ for p in model.parameters():
+ p.data.add_(p.grad, alpha=-lr)
 
  total_loss += loss.item()