File size: 1,315 Bytes
4719196 | 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 | import sys,tempfile,unittest
from pathlib import Path
import numpy as np
from PIL import Image
ROOT=Path(__file__).resolve().parents[1]
sys.path[:0]=[str(ROOT/'code'),str(ROOT/'ui')]
from train_recon import load_dat,Net
from infer_brainmu_lora import metrics,LoRALinear
from multi_gpu import shard_indices
import torch
class CoreTests(unittest.TestCase):
def test_dat_orientation(self):
x=np.zeros((41,250,400),np.uint8);x[0,0,0]=1;x[-1,-1,-1]=1
with tempfile.TemporaryDirectory() as d:
p=Path(d)/'x.dat';np.packbits(x.reshape(41,-1),axis=1,bitorder='little').tofile(p)
np.testing.assert_array_equal(load_dat(p),x[:,::-1,:])
def test_metrics_identity(self):
m=metrics(Image.new('RGB',(32,32),'gray'),Image.new('RGB',(32,32),'gray'))
self.assertAlmostEqual(m['ssim'],1);self.assertAlmostEqual(m['psnr_db'],120)
def test_shards(self):
s=shard_indices(1000,8);self.assertEqual([len(x) for x in s],[125]*8)
self.assertEqual(sorted(i for x in s for i in x),list(range(1000)))
def test_lora_zero(self):
base=torch.nn.Linear(4,3);m=LoRALinear(base,2,4,0);m.lora_A.data.zero_();m.lora_B.data.zero_();x=torch.ones(1,4)
torch.testing.assert_close(m(x),base(x))
def test_frontend(self):
self.assertEqual(tuple(Net()(torch.zeros(1,41,8,8)).shape),(1,1,8,8))
if __name__=='__main__':unittest.main()
|