(self)
| 85 | |
| 86 | @contextmanager |
| 87 | def record(self): |
| 88 | torch.cuda.reset_peak_memory_stats(self.device) |
| 89 | self._last_input_size = None |
| 90 | self._last_mem_ratio = [] |
| 91 | yield |
| 92 | self._last_memory = torch.cuda.max_memory_allocated(self.device) / 1024**3 |
| 93 | self._last_mem_ratio = sum(self._last_mem_ratio) / len(self._last_mem_ratio) |
| 94 | self._add_sample(self._last_memory, self._last_input_size, self._last_mem_ratio) |
| 95 | self.step += 1 |
| 96 | if self.step % self.update_every == 0: |
| 97 | self._max_mem_ratio = min(1.0, self._max_mem_ratio + 0.1) |
| 98 | self._fit_params() |
| 99 | |
| 100 | def _fit_params(self): |
| 101 | memory_usage = self._memory[:self._buffer_length] |
nothing calls this directly
no test coverage detected