training loop Parameters ---------- model: Any EfficientZero models target_model: Any EfficientZero models for reanalyzing replay_buffer: Any replay buffer shared_storage: Any model storage batch_storage: Any batch storage (queue)
(model, target_model, replay_buffer, shared_storage, batch_storage, config, summary_writer)
| 325 | |
| 326 | |
| 327 | def _train(model, target_model, replay_buffer, shared_storage, batch_storage, config, summary_writer): |
| 328 | """training loop |
| 329 | Parameters |
| 330 | ---------- |
| 331 | model: Any |
| 332 | EfficientZero models |
| 333 | target_model: Any |
| 334 | EfficientZero models for reanalyzing |
| 335 | replay_buffer: Any |
| 336 | replay buffer |
| 337 | shared_storage: Any |
| 338 | model storage |
| 339 | batch_storage: Any |
| 340 | batch storage (queue) |
| 341 | summary_writer: Any |
| 342 | logging for tensorboard |
| 343 | """ |
| 344 | # ---------------------------------------------------------------------------------- |
| 345 | model = model.to(config.device) |
| 346 | target_model = target_model.to(config.device) |
| 347 | |
| 348 | optimizer = optim.SGD(model.parameters(), lr=config.lr_init, momentum=config.momentum, |
| 349 | weight_decay=config.weight_decay) |
| 350 | |
| 351 | scaler = GradScaler() |
| 352 | |
| 353 | model.train() |
| 354 | target_model.eval() |
| 355 | # ---------------------------------------------------------------------------------- |
| 356 | # set augmentation tools |
| 357 | if config.use_augmentation: |
| 358 | config.set_transforms() |
| 359 | |
| 360 | # wait until collecting enough data to start |
| 361 | while not (ray.get(replay_buffer.get_total_len.remote()) >= config.start_transitions): |
| 362 | time.sleep(1) |
| 363 | pass |
| 364 | print('Begin training...') |
| 365 | # set signals for other workers |
| 366 | shared_storage.set_start_signal.remote() |
| 367 | |
| 368 | step_count = 0 |
| 369 | # Note: the interval of the current model and the target model is between x and 2x. (x = target_model_interval) |
| 370 | # recent_weights is the param of the target model |
| 371 | recent_weights = model.get_weights() |
| 372 | |
| 373 | # while loop |
| 374 | while step_count < config.training_steps + config.last_steps: |
| 375 | # remove data if the replay buffer is full. (more data settings) |
| 376 | if step_count % 1000 == 0: |
| 377 | replay_buffer.remove_to_fit.remote() |
| 378 | |
| 379 | # obtain a batch |
| 380 | batch = batch_storage.pop() |
| 381 | if batch is None: |
| 382 | time.sleep(0.3) |
| 383 | continue |
| 384 | shared_storage.incr_counter.remote() |
no test coverage detected