File size: 1,520 Bytes
9157432 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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():
# find most free gpu
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)
# device
using_cuda = len(gpu_ids) > 0 and torch.cuda.is_available()
# logger.info("Let's use %d GPUs!" % len(gpu_ids))
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()
# print(p)
return res |