Feat(PythonGC): Do garbage collection manually (#326)

* feat:add gc control

* feat:add gc control

* feat:add gc control

* feat:add gc

* re-lint
pull/338/head^2
jiaxingli 2023-09-22 13:52:25 +08:00 committed by GitHub
parent 3b0eff0c8a
commit f5337f6e02
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 0 deletions

View File

@ -2,6 +2,7 @@
# -*- encoding: utf-8 -*-
import argparse
import gc
import os
from pathlib import Path
from typing import Dict, Union
@ -446,6 +447,8 @@ def initialize_distributed_env(
master_port (str): The master port for distributed training. 8888 by default.
seed (int, optional): Specified random seed for every process. 1024 by default.
"""
# close automatic garbage collection
gc.disable()
torch.cuda.empty_cache()

View File

@ -1,6 +1,7 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import gc
import math
import socket
@ -41,6 +42,8 @@ def empty_cache_and_diag(batch_count, interval=50):
bench_net()
# do empty_cache after the bench
torch.cuda.empty_cache()
# do garbage collection
gc.collect()
def benchmark_forward(