| import torch |
| import numpy as np |
| import random |
| import pynvml |
| import logging |
|
|
|
|
| logger = logging.getLogger('MMSA') |
|
|
|
|
| def dict_to_str(src_dict): |
| dst_str = "" |
| for key in src_dict.keys(): |
| dst_str += " %s: %.4f " %(key, src_dict[key]) |
| return dst_str |
|
|
| def setup_seed(seed): |
| torch.manual_seed(seed) |
| np.random.seed(seed) |
| random.seed(seed) |
| torch.backends.cudnn.benchmark = False |
| torch.backends.cudnn.deterministic = True |
|
|
| def assign_gpu(gpu_ids, memory_limit=1e16): |
| if len(gpu_ids) == 0 and torch.cuda.is_available(): |
| |
| pynvml.nvmlInit() |
| n_gpus = pynvml.nvmlDeviceGetCount() |
| dst_gpu_id, min_mem_used = 0, memory_limit |
| for g_id in range(n_gpus): |
| handle = pynvml.nvmlDeviceGetHandleByIndex(g_id) |
| meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle) |
| mem_used = meminfo.used |
| if mem_used < min_mem_used: |
| min_mem_used = mem_used |
| dst_gpu_id = g_id |
| logger.info(f'Found gpu {dst_gpu_id}, used memory {min_mem_used}.') |
| gpu_ids.append(dst_gpu_id) |
| |
| using_cuda = len(gpu_ids) > 0 and torch.cuda.is_available() |
| |
| device = torch.device('cuda:%d' % int(gpu_ids[0]) if using_cuda else 'cpu') |
| return device |
|
|
| def count_parameters(model): |
| res = 0 |
| for p in model.parameters(): |
| if p.requires_grad: |
| res += p.numel() |
| |
| return res |
