()
| 236 | |
| 237 | |
| 238 | def GetArgumentParser(): |
| 239 | parser = argparse.ArgumentParser(description="LSTM benchmark.") |
| 240 | |
| 241 | parser.add_argument( |
| 242 | "--hidden_dim", |
| 243 | type=int, |
| 244 | default=800, |
| 245 | help="Hidden dimension", |
| 246 | ) |
| 247 | parser.add_argument( |
| 248 | "--input_dim", |
| 249 | type=int, |
| 250 | default=40, |
| 251 | help="Input dimension", |
| 252 | ) |
| 253 | parser.add_argument( |
| 254 | "--batch_size", |
| 255 | type=int, |
| 256 | default=128, |
| 257 | help="The batch size." |
| 258 | ) |
| 259 | parser.add_argument( |
| 260 | "--seq_length", |
| 261 | type=int, |
| 262 | default=20, |
| 263 | help="Max sequence length" |
| 264 | ) |
| 265 | parser.add_argument( |
| 266 | "--data_size", |
| 267 | type=int, |
| 268 | default=1000000, |
| 269 | help="Number of data points to generate" |
| 270 | ) |
| 271 | parser.add_argument( |
| 272 | "--iters_to_report", |
| 273 | type=int, |
| 274 | default=20, |
| 275 | help="Number of iteration to report progress" |
| 276 | ) |
| 277 | parser.add_argument( |
| 278 | "--gpu", |
| 279 | action="store_true", |
| 280 | help="Run all on GPU", |
| 281 | ) |
| 282 | parser.add_argument( |
| 283 | "--implementation", |
| 284 | type=str, |
| 285 | default="own", |
| 286 | help="'cudnn', 'own', 'static' or 'static_dag'", |
| 287 | ) |
| 288 | parser.add_argument( |
| 289 | "--fixed_shape", |
| 290 | action="store_true", |
| 291 | help=("Whether to randomize shape of input batches. " |
| 292 | "Static RNN requires fixed shape"), |
| 293 | ) |
| 294 | parser.add_argument( |
| 295 | "--memory_optimization", |
no outgoing calls
no test coverage detected
searching dependent graphs…