File size: 2,705 Bytes
0ea5594 | 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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 | import numpy as np
import torch
DTYPE = torch.float
DEVICE = "cuda:0"
def to_np(x):
if torch.is_tensor(x):
x = x.detach().cpu().numpy()
return x
def to_torch(x, dtype=None, device=None):
dtype = dtype or DTYPE
device = device or DEVICE
if type(x) is dict:
return {k: to_torch(v, dtype, device) for k, v in x.items()}
elif torch.is_tensor(x):
return x.to(device).type(dtype)
return torch.tensor(x, dtype=dtype, device=device)
def to_device(x, device=DEVICE):
if torch.is_tensor(x):
return x.to(device)
elif type(x) is dict:
return {k: to_device(v, device) for k, v in x.items()}
else:
raise RuntimeError(f"Unrecognized type in `to_device`: {type(x)}")
def batchify(batch):
fn = lambda x: to_torch(x[None])
batched_vals = []
for field in batch._fields:
val = getattr(batch, field)
val = apply_dict(fn, val) if type(val) is dict else fn(val)
batched_vals.append(val)
return type(batch)(*batched_vals)
def apply_dict(fn, d, *args, **kwargs):
return {k: fn(v, *args, **kwargs) for k, v in d.items()}
def normalize(x):
x = x - x.min()
x = x / x.max()
return x
def to_img(x):
normalized = normalize(x)
array = to_np(normalized)
array = np.transpose(array, (1, 2, 0))
return (array * 255).astype(np.uint8)
def set_device(device):
DEVICE = device
if "cuda" in device:
torch.set_default_tensor_type(torch.cuda.FloatTensor)
def batch_to_device(batch, device="cuda:0"):
vals = [to_device(getattr(batch, field), device) for field in batch._fields]
return type(batch)(*vals)
def _to_str(num):
if num >= 1e6:
return f"{(num/1e6):.2f} M"
else:
return f"{(num/1e3):.2f} k"
def param_to_module(param):
module_name = param[::-1].split(".", maxsplit=1)[-1][::-1]
return module_name
def report_parameters(model, topk=10):
counts = {k: p.numel() for k, p in model.named_parameters()}
n_parameters = sum(counts.values())
print(f"[ utils/arrays ] Total parameters: {_to_str(n_parameters)}")
modules = dict(model.named_modules())
sorted_keys = sorted(counts, key=lambda x: -counts[x])
max_length = max([len(k) for k in sorted_keys])
for i in range(topk):
key = sorted_keys[i]
count = counts[key]
module = param_to_module(key)
print(" " * 8, f"{key:10}: {_to_str(count)} | {modules[module]}")
remaining_parameters = sum([counts[k] for k in sorted_keys[topk:]])
print(
" " * 8,
f"... and {len(counts)-topk} others accounting for {_to_str(remaining_parameters)} parameters",
)
return n_parameters
|