|
import torch |
|
import numpy as np |
|
import os, sys |
|
import time |
|
|
|
class LyraChecker: |
|
def __init__(self, dir_data, tol): |
|
self.dir_data = dir_data |
|
self.tol = tol |
|
|
|
def cmp(self, fpath1, fpath2="", tol=0): |
|
tolbk = self.tol |
|
if tol != 0: |
|
self.tol = tol |
|
if fpath2 == "": |
|
fpath2 = fpath1 |
|
fpath1 += "_1" |
|
fpath2 += "_2" |
|
v1 = self.get_npy(fpath1) |
|
v2 = self.get_npy(fpath2) |
|
name = fpath1 |
|
if ".npy" in fpath1: |
|
name = ".".join(os.path.basename(fpath1).split(".")[:-1]) |
|
self._cmp_inner(v1, v2, name) |
|
self.tol = tolbk |
|
|
|
def _cmp_inner(self, v1, v2, name): |
|
print(v1.shape, v2.shape) |
|
if v1.shape != v2.shape: |
|
if v1.shape[1] == v2.shape[1]: |
|
v2 = v2.reshape([v2.shape[0], v2.shape[1], -1]) |
|
else: |
|
v2 = torch.tensor(v2).permute(0, 3, 1, 2).numpy() |
|
print(v1.shape, v2.shape) |
|
self._check_data(name, v1, v2) |
|
print(np.size(v1)) |
|
|
|
def _check_data(self, stage, x_out, x_gt): |
|
print(f"========== {stage} =============") |
|
print(x_out.shape, x_gt.shape) |
|
if np.allclose(x_gt, x_out, atol=self.tol): |
|
print(f"[OK] At {stage}, tol: {self.tol}") |
|
else: |
|
diff_cnt = np.count_nonzero(np.abs(x_gt - x_out)>self.tol) |
|
print(f"[FAIL]At {stage}, not aligned. tol: {self.tol}") |
|
print(" [INFO]Max diff: ", np.max(np.abs(x_gt - x_out))) |
|
print(" [INFO]Diff count: ", diff_cnt, ", ratio: ", round(diff_cnt/np.size(x_out), 2)) |
|
print(f">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>") |
|
|
|
|
|
def cmp_query(self, fpath1, fpath2): |
|
v1 = np.load(os.path.join(self.dir_data, fpath1)) |
|
vk = np.load(os.path.join(self.dir_data, fpath1).replace("query", "key")) |
|
vv = np.load(os.path.join(self.dir_data, fpath1).replace("query", "value")) |
|
|
|
v2 = np.load(os.path.join(self.dir_data, fpath2)) |
|
|
|
q2 = v2[:,:,0,:,:].transpose([0,2,1,3]) |
|
|
|
self.check_data("query", v1, q2) |
|
|
|
k2 = v2[:,:,1,:,:].transpose([0,2,1,3]) |
|
self.check_data("key", vk, k2) |
|
vv2 = v2[:,:,2,:,:].transpose([0,2,1,3]) |
|
|
|
self.check_data("value", vv, vv2) |
|
|
|
def _get_data_fpath(self, fname): |
|
fpath = os.path.join(self.dir_data, fname) |
|
if not fpath.endswith(".npy"): |
|
fpath += ".npy" |
|
return fpath |
|
|
|
def get_npy(self, fname): |
|
fpath = self._get_data_fpath(fname) |
|
return np.load(fpath) |
|
|
|
|
|
|
|
|
|
class MkDataHelper: |
|
def __init__(self, data_dir="/data/home/kiokaxiao/data"): |
|
self.data_dir = data_dir |
|
|
|
def mkdata(self, subdir, name, shape, dtype=torch.float16): |
|
outdir = os.path.join(self.data_dir, subdir) |
|
os.makedirs(outdir, exist_ok=True) |
|
fpath = os.path.join(outdir, name+".npy") |
|
data = torch.randn(shape, dtype=torch.float16) |
|
np.save(fpath, data.to(dtype).numpy()) |
|
return data |
|
|
|
def gen_out_with_func(self, func, inputs): |
|
output = func(inputs) |
|
return output |
|
|
|
def savedata(self, subdir, name, data): |
|
outdir = os.path.join(self.data_dir, subdir) |
|
os.makedirs(outdir, exist_ok=True) |
|
fpath = os.path.join(outdir, name+".npy") |
|
np.save(fpath, data.cpu().numpy()) |
|
|
|
|
|
class TorchSaver: |
|
def __init__(self, data_dir): |
|
self.data_dir = data_dir |
|
os.makedirs(self.data_dir, exist_ok=True) |
|
self.is_save = True |
|
|
|
def save_v(self, name, v): |
|
if not self.is_save: |
|
return |
|
fpath = os.path.join(self.data_dir, name+"_1.npy") |
|
np.save(fpath, v.detach().cpu().numpy()) |
|
|
|
def save_v2(self, name, v): |
|
if not self.is_save: |
|
return |
|
fpath = os.path.join(self.data_dir, name+"_1.npy") |
|
np.save(fpath, v.detach().cpu().numpy()) |
|
|
|
def timer_annoc(funct): |
|
def inner(*args,**kwargs): |
|
start = time.perf_counter() |
|
res = funct(*args,**kwargs) |
|
torch.cuda.synchronize() |
|
end = time.perf_counter() |
|
print("torch cost: ", end-start) |
|
return res |
|
return inner |
|
|
|
def get_mem_use(): |
|
f = os.popen("nvidia-smi | grep MiB" ) |
|
line = f.read().strip() |
|
while " " in line: |
|
line = line.replace(" ", " ") |
|
memuse = line.split(" ")[8] |
|
return memuse |
|
|
|
if __name__ == "__main__": |
|
dir_data = sys.argv[1] |
|
fname_v1 = sys.argv[2] |
|
fname_v2 = sys.argv[3] |
|
tol = 0.01 |
|
if len(sys.argv) > 4: |
|
tol = float(sys.argv[4]) |
|
checker = LyraChecker(dir_data, tol) |
|
checker.cmp(fname_v1, fname_v2) |
