forked from AntixK/PyTorch-VAE
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_gvae.py
40 lines (26 loc) · 920 Bytes
/
test_gvae.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
33
34
35
36
37
38
39
40
import torch
import unittest
from models import GammaVAE
from torchsummary import summary
class TestGammaVAE(unittest.TestCase):
def setUp(self) -> None:
# self.model2 = VAE(3, 10)
self.model = GammaVAE(3, 10)
def test_summary(self):
print(summary(self.model, (3, 64, 64), device='cpu'))
# print(summary(self.model2, (3, 64, 64), device='cpu'))
def test_forward(self):
x = torch.randn(16, 3, 64, 64)
y = self.model(x)
print("Model Output size:", y[0].size())
# print("Model2 Output size:", self.model2(x)[0].size())
def test_loss(self):
x = torch.randn(16, 3, 64, 64)
result = self.model(x)
loss = self.model.loss_function(*result, M_N = 0.005)
print(loss)
def test_sample(self):
self.model.cuda()
y = self.model.sample(144, 0)
if __name__ == '__main__':
unittest.main()