| 151 | |
| 152 | |
| 153 | class ThroughputTimer: |
| 154 | |
| 155 | def __init__( |
| 156 | self, |
| 157 | batch_size, |
| 158 | start_step=2, |
| 159 | steps_per_output=50, |
| 160 | monitor_memory=False, |
| 161 | logging_fn=None, |
| 162 | ): |
| 163 | from deepspeed.utils import logger |
| 164 | self.start_time = 0 |
| 165 | self.end_time = 0 |
| 166 | self.started = False |
| 167 | self.batch_size = 1 if batch_size is None else batch_size |
| 168 | self.start_step = start_step |
| 169 | self.epoch_count = 0 |
| 170 | self.micro_step_count = 0 |
| 171 | self.global_step_count = 0 |
| 172 | self.total_elapsed_time = 0 |
| 173 | self.step_elapsed_time = 0 |
| 174 | self.steps_per_output = steps_per_output |
| 175 | self.monitor_memory = monitor_memory |
| 176 | self.logging = logging_fn |
| 177 | if self.logging is None: |
| 178 | self.logging = logger.info |
| 179 | self.initialized = False |
| 180 | |
| 181 | if self.monitor_memory and not PSUTILS_INSTALLED: |
| 182 | raise ImportError("Unable to import 'psutils', please install package") |
| 183 | |
| 184 | def update_epoch_count(self): |
| 185 | self.epoch_count += 1 |
| 186 | self.micro_step_count = 0 |
| 187 | |
| 188 | def _init_timer(self): |
| 189 | self.initialized = True |
| 190 | |
| 191 | def start(self): |
| 192 | self._init_timer() |
| 193 | self.started = True |
| 194 | if self.global_step_count >= self.start_step: |
| 195 | get_accelerator().synchronize() |
| 196 | self.start_time = time.time() |
| 197 | |
| 198 | def stop(self, global_step=False, report_speed=True): |
| 199 | if not self.started: |
| 200 | return |
| 201 | self.started = False |
| 202 | self.micro_step_count += 1 |
| 203 | if global_step: |
| 204 | self.global_step_count += 1 |
| 205 | |
| 206 | if self.start_time > 0: |
| 207 | get_accelerator().synchronize() |
| 208 | self.end_time = time.time() |
| 209 | duration = self.end_time - self.start_time |
| 210 | self.total_elapsed_time += duration |