forked from tianyic/only_train_once_personal_footprint
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_diffmodel_cifar.py
32 lines (25 loc) · 1.28 KB
/
test_diffmodel_cifar.py
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
import torch
from only_train_once import OTO
from backends import DiffModelCIFAR
import unittest
import os
OUT_DIR = './cache'
class TestDiffModelCIFAR(unittest.TestCase):
def test_sanity(self, dummy_input=torch.rand(1, 4, 32, 32)):
model = DiffModelCIFAR()
oto = OTO(model, dummy_input)
# The layout rendering of DiffUnet is out-of-time via graphviz
# oto.visualize(view=False, out_dir=OUT_DIR, display_params=True)
oto.random_set_zero_groups()
oto.construct_subnet(out_dir=OUT_DIR)
full_model = torch.load(oto.full_group_sparse_model_path)
compressed_model = torch.load(oto.compressed_model_path)
full_output = full_model(dummy_input)
compressed_output = compressed_model(dummy_input)
max_output_diff = torch.max(torch.abs(full_output - compressed_output))
print("Maximum output difference " + str(max_output_diff.item()))
# self.assertLessEqual(max_output_diff, 1e-4)
full_model_size = os.stat(oto.full_group_sparse_model_path)
compressed_model_size = os.stat(oto.compressed_model_path)
print("Size of full model : ", full_model_size.st_size / (1024 ** 3), "GBs")
print("Size of compress model : ", compressed_model_size.st_size / (1024 ** 3), "GBs")