class ActnormTest(tf.test.TestCase):
def setUp(self):
super(ActnormTest, self).setUp()
self.actnorm = Actnorm()
self.actnorm.build((None, 16, 16, 4))
def testActnormInitializeOutputShape(self):
self.assertFalse(self.actnorm.initialized)
x = tf.random.normal([1024, 16, 16, 4])
z, ldj = self.actnorm(x)
self.assertTrue(self.actnorm.initialized)
self.assertShapeEqual(np.zeros(x.shape), z)
self.assertShapeEqual(np.zeros(x.shape[0:1]), ldj)
def testActnormOutput(self):
x = tf.random.normal([1024, 16, 16, 4])
self.actnorm(x)
self.assertShapeEqual(np.zeros([1, 1, 1, 4]),
self.actnorm.logs.value())
self.assertShapeEqual(np.zeros([1, 1, 1, 4]),
self.actnorm.bias.value())
self.assertAllClose(
self.actnorm.logs.value(),
tf.math.log(self.actnorm.scale / tf.ones([1, 1, 1, 4])
# / self.actnorm.logscale_factor
),
# * self.actnorm.logscale_factor,
rtol=1e-8,
atol=1e-1,
)
self.assertAllClose(self.actnorm.bias.value(),
tf.zeros([1, 1, 1, 4]),
rtol=1e-8,
atol=1e-1)
class ActnormScaleTest(tf.test.TestCase):
def setUp(self):
super(ActnormScaleTest, self).setUp()
self.scale = 2.0
self.actnorm = Actnorm(scale=self.scale)
self.actnorm.build((None, 16, 16, 4))
def testActnormInitializeOutputShape(self):
x = tf.random.normal([1024, 16, 16, 4])
z, ldj = self.actnorm(x)
self.assertShapeEqual(np.zeros(x.shape), z)
self.assertShapeEqual(np.zeros(x.shape[0:1]), ldj)
def __init__(
self,
K: int = 5,
L: int = 3,
resblk_kwargs: Dict = None,
conditional: bool = False,
):
super(BasicGlow, self).__init__()
if resblk_kwargs is None:
resblk_kwargs = {"num_block": 3, "units_factor": 6}
self.resblk_kwargs = resblk_kwargs
self.K = K
self.L = L
layers = []
layers.append(LogitifyImage())
for l in range(self.L):
if l == 0:
layers.append(Squeeze(with_zaux=False))
else:
layers.append(Squeeze(with_zaux=True))
fml = []
for k in range(self.K):
fml.append(Actnorm())
fml.append(Inv1x1Conv())
fml.append(
AffineCoupling(scale_shift_net=ShallowResNet(
**self.resblk_kwargs)))
layers.append(FlowModule(fml))
if l == 0:
layers.append(FactorOut(conditional=conditional))
elif l != self.L - 1:
layers.append(
FactorOut(with_zaux=True, conditional=conditional))
self.flows = layers
class ActnormInvTest(tf.test.TestCase):
def setUp(self):
super(ActnormInvTest, self).setUp()
self.actnorm = Actnorm()
self.actnorm.build((None, 16, 16, 4))
def testActnormInitializeOutputShape(self):
x = tf.random.normal([1024, 16, 16, 4])
z, ldj = self.actnorm(x)
self.assertShapeEqual(np.zeros(x.shape), z)
self.assertShapeEqual(np.zeros(x.shape[0:1]), ldj)
def testActnormInvOutput(self):
x = tf.random.normal([1024, 16, 16, 4])
self.actnorm(x)
x = tf.random.normal([1024, 16, 16, 4])
z, ldj = self.actnorm(x)
rev_x, ildj = self.actnorm(z, inverse=True)
self.assertAllClose(ldj + ildj,
tf.zeros(x.shape[0:1]),
rtol=1e-8,
atol=1e-8)
self.assertAllClose(x, rev_x, rtol=1e-8, atol=1e-5)
def __init__(self, K: int = 5, L: int = 1, resblk_kwargs: Dict = None):
super().__init__()
if resblk_kwargs is None:
resblk_kwargs = {"num_block": 3, "units_factor": 6}
self.resblk_kwargs = resblk_kwargs
self.K = K
self.L = L
layers = []
layers.append(LogitifyImage())
for _ in range(3):
layers.append(Squeeze())
for _ in range(5):
layers.append(Actnorm())
layers.append(Inv1x1Conv())
layers.append(
AffineCoupling(scale_shift_net=ShallowResNet(
**self.resblk_kwargs)))
self.flows = layers
def setUp(self):
super(ActnormTest, self).setUp()
self.actnorm = Actnorm()
self.actnorm.build((None, 16, 16, 4))
def setUp(self):
super(ActnormScaleTest, self).setUp()
self.scale = 2.0
self.actnorm = Actnorm(scale=self.scale)
self.actnorm.build((None, 16, 16, 4))