| 197 | |
| 198 | |
| 199 | class ThroughputTimer: |
| 200 | |
| 201 | def __init__(self, config, batch_size, start_step=2, steps_per_output=None, monitor_memory=False, logging_fn=None): |
| 202 | from deepspeed.utils import logger |
| 203 | self.config = config |
| 204 | self.start_time = 0 |
| 205 | self.end_time = 0 |
| 206 | self.started = False |
| 207 | self.batch_size = 1 if batch_size is None else batch_size |
| 208 | self.start_step = start_step |
| 209 | self.epoch_count = 0 |
| 210 | self.micro_step_count = 0 |
| 211 | self.global_step_count = 0 |
| 212 | self.total_elapsed_time = 0 |
| 213 | self.step_elapsed_time = 0 |
| 214 | self.steps_per_output = steps_per_output |
| 215 | self.monitor_memory = monitor_memory |
| 216 | self.logging = logging_fn |
| 217 | if self.logging is None: |
| 218 | self.logging = logger.info |
| 219 | self.initialized = False |
| 220 | |
| 221 | if self.monitor_memory and not PSUTILS_INSTALLED: |
| 222 | raise ImportError("Unable to import 'psutils', please install package") |
| 223 | |
| 224 | def update_epoch_count(self): |
| 225 | self.epoch_count += 1 |
| 226 | self.micro_step_count = 0 |
| 227 | |
| 228 | def _init_timer(self): |
| 229 | self.initialized = True |
| 230 | |
| 231 | def start(self): |
| 232 | if not self.config.enabled: |
| 233 | return |
| 234 | self._init_timer() |
| 235 | self.started = True |
| 236 | if self.global_step_count >= self.start_step: |
| 237 | if self.config.synchronized: |
| 238 | get_accelerator().synchronize() |
| 239 | self.start_time = time.time() |
| 240 | |
| 241 | def _is_report_boundary(self): |
| 242 | if self.steps_per_output is None: |
| 243 | return False |
| 244 | return self.global_step_count % self.steps_per_output == 0 |
| 245 | |
| 246 | def stop(self, global_step=False, report_speed=True): |
| 247 | if not self.config.enabled or not self.started: |
| 248 | return |
| 249 | self.started = False |
| 250 | self.micro_step_count += 1 |
| 251 | if global_step: |
| 252 | self.global_step_count += 1 |
| 253 | |
| 254 | if self.start_time > 0: |
| 255 | if self.config.synchronized: |
| 256 | get_accelerator().synchronize() |