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 LogitifyImageTest(tf.test.TestCase):
def setUp(self):
super().setUp()
self.li = LogitifyImage()
self.li.build([None, 32, 32, 1])
def testLogitifyImageOutputShape(self):
x = tf.nn.tanh(tf.random.normal([1024, 32, 32, 1]))
z, ldj = self.li(x, inverse=False)
self.assertShapeEqual(np.zeros(x.shape), z)
self.assertShapeEqual(np.zeros(x.shape[0:1]), ldj)
def testLogitifyImageOutput(self):
# x's range is [0, 1]
x = tf.nn.sigmoid(tf.random.normal([1024, 32, 32, 1]))
z, ldj = self.li(x, inverse=False)
rev_x, ildj = self.li(z, inverse=True)
self.assertAllClose(x, rev_x, rtol=8e-1, atol=1e-2)
self.assertAllClose(ldj + ildj, tf.zeros([1024]))
def __init__(self, L=3):
super(SqueezeFactorOut, self).__init__()
self.flows = [
LogitifyImage(),
Squeeze(),
FactorOut(),
Squeeze(with_zaux=True),
FactorOut(with_zaux=True),
Squeeze(with_zaux=True),
FactorOut(with_zaux=True),
]
def _main():
layer = LogitifyImage() # BasicGlow()
x = tf.keras.Input((32, 32, 1))
y = layer(x, training=True)
model = tf.keras.Model(x, y)
train, test = tf.keras.datasets.mnist.load_data()
train_image = train[0] / 255.0
train_image = train_image[..., tf.newaxis]
# forward -> inverse
train_image = train_image[0:12]
forward, ldj = layer.forward(train_image)
inverse, ildj = layer.inverse(forward)
print(ldj)
print(ildj)
print(ldj + ildj)
print(tf.reduce_mean(ldj + ildj))
print(tf.reduce_mean(train_image - inverse))
train_image = inverse
print(train_image.shape)
import matplotlib.pyplot as plt
fig = plt.figure(figsize=(18, 18))
for i in range(9):
img = tf.squeeze(train_image[i])
fig.add_subplot(3, 3, i + 1)
plt.title(train[1][i])
plt.tick_params(bottom=False,
left=False,
labelbottom=False,
labelleft=False)
plt.imshow(img, cmap="gray_r")
plt.show(block=True)
model.summary()
return model
def __init__(self, hparams: Dict):
super().__init__()
self.model_params = hparams["model"]
K = self.model_params["K"]
L = self.model_params["L"]
conditional = self.model_params["conditional"]
flows = []
flows.append(LogitifyImage())
for layer in range(L):
# Squeezing Layer
if layer == 0:
flows.append(Squeezing(with_zaux=False))
else:
flows.append(Squeezing(with_zaux=True))
fml = []
# build flow module layer
for k in range(K):
fml.append(Actnorm())
fml.append(Inv1x1Conv())
scale_shift_net = ShallowResNet(width=self.model_params["hidden_width"])
fml.append(
ConditionalAffineCoupling(
mask_type=AffineCouplingMask.ChannelWise,
scale_shift_net=scale_shift_net,
)
)
flows.append(ConditionalFlowModule(fml))
# Factor Out Layer
if layer == 0:
flows.append(FactorOut(conditional=conditional))
elif layer != L - 1:
flows.append(FactorOut(with_zaux=True, conditional=conditional))
self.flows = flows
self.cond_conf = hparams["images"]["cond"]
self.cond_net = CondNet(
[
self.cond_conf["height"],
self.cond_conf["width"],
self.cond_conf["channel"],
]
)
self.cond_net.summary()
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 __init__(self, hparams: Dict):
super().__init__()
self.model_params = hparams["model"]
K = self.model_params["K"]
L = self.model_params["L"]
conditional = self.model_params["conditional"]
flows = []
flows.append(LogitifyImage())
for layer in range(L):
# Squeeze Layer
if layer == 0:
flows.append(Squeeze(with_zaux=False))
else:
flows.append(Squeeze(with_zaux=True))
fml = []
# build flow module layer
for k in range(K):
fml.append(Actnorm())
fml.append(Inv1x1Conv())
fml.append(
AffineCoupling(
mask_type=AffineCouplingMask.ChannelWise,
scale_shift_net_template=lambda x: ShallowResNet(
x, width=self.model_params["hidden_width"]),
))
flows.append(FlowModule(fml))
# Factor Out Layer
if layer == 0:
flows.append(FactorOut(conditional=conditional))
elif layer != L - 1:
flows.append(FactorOut(with_zaux=True,
conditional=conditional))
self.flows = flows
def setUp(self):
super().setUp()
self.li = LogitifyImage()
self.li.build([None, 32, 32, 1])