| 319 | } |
| 320 | |
| 321 | void memory_profiler::on_setup_end(model* m) |
| 322 | { |
| 323 | auto comm = m->get_comm(); |
| 324 | bool should_print = comm->am_trainer_master(); |
| 325 | |
| 326 | // Post-setup printout of layer accounting |
| 327 | if (should_print) { |
| 328 | std::cout << "MEM: Expected memory usage by layer (in descending order):" |
| 329 | << std::endl; |
| 330 | report_mem_usage(m); |
| 331 | } |
| 332 | |
| 333 | // Print total used memory |
| 334 | m_step0_usage = m_setup_end_usage = get_used_gpu_memory(); |
| 335 | if (m_setup_end_usage > m_initial_memory_usage && should_print) { |
| 336 | std::cout << "MEM: Total actual memory usage after setup: " |
| 337 | << (m_setup_end_usage - m_initial_memory_usage) / 1048576.0 |
| 338 | << " MiB." << std::endl; |
| 339 | } |
| 340 | m_current_step = 0; |
| 341 | } |
| 342 | |
| 343 | void memory_profiler::first_step_accounting(model* m, const std::string& msg) |
| 344 | { |
no test coverage detected