Spaces:
Running on Zero
Running on Zero
File size: 2,339 Bytes
21e9df7 | 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 105 106 107 108 109 110 111 112 113 | import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import wraps
from einops import rearrange, reduce
import math
def once(fn):
called = False
@wraps(fn)
def inner(x):
nonlocal called
if called:
return
called = True
return fn(x)
return inner
print_once = once(print)
# functions
def exists(val):
return val is not None
def default(val, d):
return val if exists(val) else d
def Sequential(*modules):
return nn.Sequential(*filter(exists, modules))
# tensor functions
def log(t, eps=1e-20):
return torch.log(t.clamp(min=eps))
def l2norm(t):
return F.normalize(t, p=2, dim=-1)
def matrix_diag(t):
device = t.device
i, j = t.shape[-2:]
num_diag_el = min(i, j)
i_range = torch.arange(i, device=device)
j_range = torch.arange(j, device=device)
diag_mask = rearrange(i_range, 'i -> i 1') == rearrange(j_range, 'j -> 1 j')
diag_el = t.masked_select(diag_mask)
return rearrange(diag_el, '(b d) -> b d', d=num_diag_el)
# 2d sinusoidal positional embedding
# simple vit paper shows it is good enough compared to learned
class LayerNorm(nn.Module):
def __init__(self, dim, scale=True):
super().__init__()
self.learned_gamma = nn.Parameter(torch.ones(dim)) if scale else None
self.register_buffer('gamma', torch.ones(dim), persistent=False)
self.register_buffer('beta', torch.zeros(dim), persistent=False)
def forward(self, x):
return F.layer_norm(x, x.shape[-1:], default(self.learned_gamma, self.gamma), self.beta)
def freeze(model):
for n, p in model.named_parameters():
p.requires_grad = False
def print_trainable_parameters(model):
trainable_params = 0
all_param = 0
for k, param in model.named_parameters():
num_params = param.numel()
# if using DS Zero 3 and the weights are initialized empty
if num_params == 0 and hasattr(param, "ds_numel"):
num_params = param.ds_numel
all_param += num_params
if param.requires_grad:
print(k)
trainable_params += num_params
print(
f"trainable params: {trainable_params:,d} || all params: {all_param:,d} || trainable%: {100 * trainable_params / all_param}"
) |