|
|
|
import torch |
|
|
|
from .hook import HOOKS, Hook |
|
|
|
|
|
@HOOKS.register_module() |
|
class EmptyCacheHook(Hook): |
|
|
|
def __init__(self, before_epoch=False, after_epoch=True, after_iter=False): |
|
self._before_epoch = before_epoch |
|
self._after_epoch = after_epoch |
|
self._after_iter = after_iter |
|
|
|
def after_iter(self, runner): |
|
if self._after_iter: |
|
torch.cuda.empty_cache() |
|
|
|
def before_epoch(self, runner): |
|
if self._before_epoch: |
|
torch.cuda.empty_cache() |
|
|
|
def after_epoch(self, runner): |
|
if self._after_epoch: |
|
torch.cuda.empty_cache() |
|
|
