stop tracking for the passed stage
(self, stage)
| 757 | peak_monitor_thread.start() |
| 758 | |
| 759 | def stop(self, stage): |
| 760 | """stop tracking for the passed stage""" |
| 761 | |
| 762 | # deal with nested calls of eval during train - simply ignore those |
| 763 | if self.cur_stage is not None and self.cur_stage != stage: |
| 764 | return |
| 765 | |
| 766 | # this sends a signal to peak_monitor_func to complete its loop |
| 767 | self.peak_monitoring = False |
| 768 | |
| 769 | # first ensure all objects get collected and their memory is freed |
| 770 | gc.collect() |
| 771 | |
| 772 | if self.torch is not None: |
| 773 | if torch.cuda.is_available(): |
| 774 | self.torch.cuda.empty_cache() |
| 775 | elif is_torch_mlu_available(): |
| 776 | self.torch.mlu.empty_cache() |
| 777 | elif is_torch_musa_available(): |
| 778 | self.torch.musa.empty_cache() |
| 779 | elif is_torch_xpu_available(): |
| 780 | self.torch.xpu.empty_cache() |
| 781 | elif is_torch_npu_available(): |
| 782 | self.torch.npu.empty_cache() |
| 783 | elif is_torch_hpu_available(): |
| 784 | # not available on hpu as it reserves all device memory for the current process |
| 785 | # self.torch.npu.empty_cache() |
| 786 | pass |
| 787 | elif is_torch_mps_available(): |
| 788 | self.torch.mps.empty_cache() |
| 789 | |
| 790 | # concepts: |
| 791 | # - alloc_delta: the difference of allocated memory between the end and the start |
| 792 | # - peaked_delta: the difference between the peak memory and the current memory |
| 793 | # in order to know how much memory the measured code consumed one needs to sum these two |
| 794 | |
| 795 | # gpu |
| 796 | if self.torch is not None: |
| 797 | if torch.cuda.is_available(): |
| 798 | self.gpu_mem_used_now = self.torch.cuda.memory_allocated() |
| 799 | self.gpu_mem_used_peak = self.torch.cuda.max_memory_allocated() |
| 800 | elif is_torch_mlu_available(): |
| 801 | self.gpu_mem_used_now = self.torch.mlu.memory_allocated() |
| 802 | self.gpu_mem_used_peak = self.torch.mlu.max_memory_allocated() |
| 803 | elif is_torch_musa_available(): |
| 804 | self.gpu_mem_used_now = self.torch.musa.memory_allocated() |
| 805 | self.gpu_mem_used_peak = self.torch.musa.max_memory_allocated() |
| 806 | elif is_torch_xpu_available(): |
| 807 | self.gpu_mem_used_now = self.torch.xpu.memory_allocated() |
| 808 | self.gpu_mem_used_peak = self.torch.xpu.max_memory_allocated() |
| 809 | elif is_torch_npu_available(): |
| 810 | self.gpu_mem_used_now = self.torch.npu.memory_allocated() |
| 811 | self.gpu_mem_used_peak = self.torch.npu.max_memory_allocated() |
| 812 | elif is_torch_hpu_available(): |
| 813 | self.gpu_mem_used_now = self.torch.hpu.memory_allocated() |
| 814 | self.gpu_mem_used_peak = self.torch.hpu.max_memory_allocated() |
| 815 | elif is_torch_mps_available(): |
| 816 | self.gpu_mem_used_now = self.torch.mps.current_allocated_memory() |
no test coverage detected