def preserve_gpu_with_id(gpu_id: int, preserve_percent: float = 0.95):
logger = log_utils.get_master_logger("preserve_gpu_with_id")
if not torch.cuda.is_available():
logger.warning("no gpu avaliable exit...")
return
try:
import cupy
device = cupy.cuda.Device(gpu_id)
avaliable_mem = device.mem_info[0] - 700 * 1024 * 1024
logger.info("{}MB memory avaliable, trying to preserve {}MB...".format(
int(avaliable_mem / 1024.0 / 1024.0),
int(avaliable_mem / 1024.0 / 1024.0 * preserve_percent)))
if avaliable_mem / 1024.0 / 1024.0 < 100:
cmd = os.popen("nvidia-smi")
outputs = cmd.read()
pid = os.getpid()
logger.warning("Avaliable memory is less than 100MB, skiping...")
logger.info("program pid: %d, current environment:\n%s", pid,
outputs)
raise MemoryPreserveError()
alloc_mem = int(avaliable_mem * preserve_percent / 4)
x = torch.empty(alloc_mem).to(torch.device("cuda:{}".format(gpu_id)))
del x
except ImportError:
logger.warning("No cupy found, memory cannot be perserved")
def __init__(
self,
root: str,
split: str = "train",
version: str = "2017",
resize: Tuple[int] = (300, 300),
augmentations: Callable[[Image, FloatTensor, LongTensor], Tuple[Image, FloatTensor, LongTensor]] = None,
keep_crowd: bool = False
):
super().__init__(resize, augmentations)
self.root = os.path.expanduser(root)
verify_str_arg(split, "split", ("train", "val"))
self.keep_crowd = keep_crowd
self.logger = log_utils.get_master_logger("CocoDetection")
# path to image folder, e.g. coco_root/train2017
self.image_folder = os.path.join(self.root, f"{split}{version}")
# path to annotation, e.g. coco_root/annotations/instances_train2017
annotation_fp = os.path.join(self.root, "annotations", f"instances_{split}{version}.json")
self.logger.info("Parsing COCO %s dataset...", split)
self._init_dataset(annotation_fp)
self.logger.info("Parsing COCO %s dataset done", split)
def get_composed_augmentations(aug_list: List[Dict[str, Any]]):
logger = log_utils.get_master_logger("get_composed_augmentations")
augmentations = list()
for aug in aug_list:
assert len(aug) == 1, "one aug each time"
name, args = aug.popitem()
aug_method = __REGISTERED_AUGS__[name]
if isinstance(args, (list, tuple)):
augmentations.append(aug_method(*args))
elif isinstance(args, dict):
augmentations.append(aug_method(**args))
else:
augmentations.append(aug_method(args))
logger.info("Using {} aug with params {}".format(name, args))
return Compose(*augmentations)
def preserve_memory(preserve_percent: float = 0.99):
logger = log_utils.get_master_logger("preserve_memory")
if not torch.cuda.is_available():
logger.warning("no gpu avaliable exit...")
return
thread_pool = list()
for i in range(torch.cuda.device_count()):
thread = threading.Thread(target=preserve_gpu_with_id,
kwargs=dict(
gpu_id=i,
preserve_percent=preserve_percent),
name="Preserving GPU {}".format(i))
logger.info("Starting to preserve GPU: {}".format(i))
thread.start()
thread_pool.append(thread)
for t in thread_pool:
t.join()
def __init__(self,
root: str,
split: str = "trainval",
version: str = "2007",
resize: Optional[Tuple[int]] = (300, 300),
augmentations: Callable[[Image, Dict[str, Any]],
Tuple[Image, Dict[str, Any]]] = None,
keep_difficult: bool = False,
make_partial: List[int] = None):
"""
Args:
root: root to voc path which contains [`VOC2007` `VOC2012`] folders
split: split of dataset, e.g. `train`, `val`, `test` and `trainval`. Warning: VOC2012 has no
test split
version: `2007` or `2012`
resize: all images will be resized to given size. If `None`, all images will not be resized
make_partial: only keep objects with given classes, w.r.t `self.CLASSES`. If `None`, all
objects will be perserved. For multitask learning which one task has 10 classes and the
other task has others classes
"""
super().__init__(resize, augmentations)
self.keep_difficult = keep_difficult
verify_str_arg(version, "version", ("2007", "2012"))
verify_str_arg(split, "split", ("train", "val", "test", "trainval"))
self.logger = log_utils.get_master_logger("VOCDetection")
self.version = version
self.split = split
# parse folders
self.root = os.path.expanduser(os.path.join(root, f"VOC{version}"))
# read split file
split_fp = os.path.join(self.root, "ImageSets", "Main", f"{split}.txt")
if not os.path.isfile(split_fp):
raise FileNotFoundError(
f"`{split_fp}` is not found, note that there is no `test.txt` for VOC-2012"
)
with open(split_fp, "r") as f:
self.file_names = [x.strip() for x in f.readlines()]
self.logger.info("Parsing VOC%s %s dataset...", version, split)
self._init_dataset(make_partial)
self.logger.info("Parsing VOC%s %s dataset done", version, split)
def get_optimizer(params: List[Dict[str, nn.Parameter]],
optimizer_cfg: Dict[str, Dict[str, Any]]) -> Optimizer:
"""
optimizer yml example
```
optimizer:
name: sgd
lr: 0.01
weight_decay: 5.0e-4
momentum: 0.9
```
"""
logger = log_utils.get_master_logger("get_optimizer")
opt_name = optimizer_cfg["name"]
optim_cfg = copy.deepcopy(optimizer_cfg)
optim_cfg.pop("name")
logger.info("Using {} optimizer with config {}".format(
opt_name, to_json_str(optim_cfg)))
optimizer = __REGISTERED_OPTIMIZERS__[opt_name](params, **optim_cfg)
return optimizer
def get_scheduler(optimizer, scheduler_cfg: Dict[str, Dict[str, Any]]):
"""
scheduler yml example
```
lr_schedule:
name: multi_step
milestones: [8000]
gamma: 0.1
```
"""
logger = log_utils.get_master_logger("get_scheduler")
scheduler_dict = deepcopy(scheduler_cfg)
if scheduler_dict is None:
logger.info("Using No lr scheduling, fallback to constant lr")
return ConstantLR(optimizer)
s_type = scheduler_dict.pop("name")
logger.info("Using {} scheduler with\n{}".format(
s_type, to_json_str(scheduler_dict)))
if "warmup_iters" in scheduler_dict:
warmup_dict = {}
warmup_dict["warmup_iters"] = scheduler_dict.pop("warmup_iters", 100)
warmup_dict["mode"] = scheduler_dict.pop("warmup_mode", "linear")
warmup_dict["gamma"] = scheduler_dict.pop("warmup_factor", 0.2)
logger.info("Using Warmup with {} iters {} gamma and {} mode".format(
warmup_dict["warmup_iters"], warmup_dict["gamma"],
warmup_dict["mode"]))
base_scheduler = __REGISTERED_SCHEDULERS__[s_type](optimizer,
**scheduler_dict)
return WarmUpLR(optimizer, base_scheduler, **warmup_dict)
return __REGISTERED_SCHEDULERS__[s_type](optimizer, **scheduler_dict)