File size: 1,291 Bytes
3d5551b | 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 | import torch
import random
import numpy as np
from typing import List
from itertools import repeat
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def grad_status(model):
return (par.requires_grad for par in model.parameters())
def lmap(f, x):
"""list(map(f, x))"""
return list(map(f, x))
def assert_all_frozen(model):
model_grads: List[bool] = list(grad_status(model))
n_require_grad = sum(lmap(int, model_grads))
npars = len(model_grads)
assert not any(
model_grads
), f"{n_require_grad / npars:.1%} of {npars} weights require grad"
def split_dense_inputs(model_input: dict, chunk_size: int):
assert len(model_input) == 1
arg_key = list(model_input.keys())[0]
arg_val = model_input[arg_key]
keys = list(arg_val.keys())
chunked_tensors = [arg_val[k].split(chunk_size, dim=0) for k in keys]
chunked_arg_val = [
dict(zip(kk, tt)) for kk, tt in zip(repeat(keys), zip(*chunked_tensors))
]
return [{arg_key: c} for c in chunked_arg_val]
def get_dense_rep(x):
if x.q_reps is None:
return x.p_reps
else:
return x.q_reps
|