def test_saveload():
p = source.Gaussian([2,2])
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(1,4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:,i] = 1
b=b.reshape(1,2,2)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
f = flow.RNVP(maskList, [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),utils.ScalableTanh(4)]) for _ in range(4)],p)
p = source.Gaussian([2,2])
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(1,4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:,i] = 1
b=b.reshape(1,2,2)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
blankf = flow.RNVP(maskList, [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),utils.ScalableTanh(4)]) for _ in range(4)],p)
saveload(f,blankf)
def test_saveload():
decimal = flow.ScalingNshifting(256, -128)
p = source.MixtureDiscreteLogistic([3, 32, 32], 5, decimal, utils.roundingWidentityGradient)
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.cat([torch.zeros(3 * 32 * 16), torch.ones(3 * 32 * 16)])[torch.randperm(3 * 32 * 32)].reshape(1, 3, 32, 32)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
tList = [utils.SimpleMLPreshape([3 * 32 * 16, 200, 500, 3 * 32 * 16], [nn.ELU(), nn.ELU(), None]) for _ in range(4)]
f = flow.DiscreteNICE(maskList, tList, decimal, utils.roundingWidentityGradient, p)
p = source.MixtureDiscreteLogistic([3, 32, 32], 5, decimal, utils.roundingWidentityGradient)
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.cat([torch.zeros(3 * 32 * 16), torch.ones(3 * 32 * 16)])[torch.randperm(3 * 32 * 32)].reshape(1, 3, 32, 32)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
tList = [utils.SimpleMLPreshape([3 * 32 * 16, 200, 500, 3 * 32 * 16], [nn.ELU(), nn.ELU(), None]) for _ in range(4)]
blankf = flow.DiscreteNICE(maskList, tList, decimal, utils.roundingWidentityGradient, p)
saveload(f,blankf)
def test_bijective():
p = source.Gaussian([1, 4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.view(1, 1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
depth = 4
t = flow.TEBD(2, length, layers, depth, p)
bijective(t)
def symmetryMERAInit(L,
d,
nlayers,
nmlp,
nhidden,
nrepeat,
symmetryList,
device,
dtype,
name=None,
channel=1,
depthMERA=None):
s = source.Gaussian([channel] + [L] * d)
depth = int(math.log(L, 2)) * nrepeat * 2
coreSize = 4 * channel
MaskList = []
for _ in range(depth):
masklist = []
for n in range(nlayers):
if n % 2 == 0:
b = torch.zeros(1, coreSize)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.view(1, channel, 2, 2)
else:
b = 1 - b
masklist.append(b)
masklist = torch.cat(masklist, 0).to(torch.float32)
MaskList.append(masklist)
dimList = [coreSize]
for _ in range(nmlp):
dimList.append(nhidden)
dimList.append(coreSize)
layers = [
flow.RNVP(MaskList[n], [
utils.SimpleMLPreshape(dimList, [nn.ELU()
for _ in range(nmlp)] + [None])
for _ in range(nlayers)
], [
utils.SimpleMLPreshape(dimList, [nn.ELU() for _ in range(nmlp)] +
[utils.ScalableTanh(coreSize)])
for _ in range(nlayers)
]) for n in range(depth)
]
f = flow.MERA(2, L, layers, nrepeat, depth=depthMERA, prior=s)
if symmetryList is not None:
f = Symmetrized(f, symmetryList, name=name)
f.to(device=device, dtype=dtype)
return f
def test_bijective():
p = source.Gaussian([4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.view(1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
repeat = 2
t = flow.MERA(2, length, layers, repeat, p)
def op(x):
return -x
sym = [op]
m = train.Symmetrized(t, sym)
z = m.prior.sample(100)
xz1, _ = m.inverse(z)
xz2, _ = m.inverse(z)
p1 = m.logProbability(xz1)
p2 = m.logProbability(xz2)
z1, _ = m.forward(xz1)
xz1p, _ = m.inverse(z1)
assert ((xz1 == xz2).sum() + (xz1 == -xz2).sum()) == 100 * 4 * 4
assert_array_almost_equal(p1.detach().numpy(),
p2.detach().numpy(),
decimal=5)
assert_array_almost_equal(np.fabs(xz1.detach().numpy()),
np.fabs(xz1p.detach().numpy()))
def test_grad():
decimal = flow.ScalingNshifting(256, -128)
p = source.MixtureDiscreteLogistic([3, 32, 32], 5, decimal, utils.roundingWidentityGradient)
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.cat([torch.zeros(3 * 32 * 16), torch.ones(3 * 32 * 16)])[torch.randperm(3 * 32 * 32)].reshape(1, 3, 32, 32)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
tList = [utils.SimpleMLPreshape([3 * 32 * 16, 200, 500, 3 * 32 * 16], [nn.ELU(), nn.ELU(), None]) for _ in range(4)]
f = flow.DiscreteNICE(maskList, tList, decimal, utils.roundingWidentityGradient, p)
fcopy = deepcopy(f)
fcopy.rounding = torch.round
field = p.sample(100).detach()
cfield = deepcopy(field).requires_grad_()
field.requires_grad_()
xfield, _ = f.inverse(field)
xcfield, _ = fcopy.inverse(cfield)
L = xfield.sum()
Lc = xcfield.sum()
L.backward()
Lc.backward()
ou = [term for term in f.parameters()]
ouc = [term for term in fcopy.parameters()]
assert not np.all(ou[-1].grad.detach().numpy() == ouc[-1].grad.detach().numpy())
def test_bijective():
p = source.Gaussian([3, 2, 2])
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
f = flow.OnebyonePlusRNVP(maskList, [
utils.SimpleMLPreshape([12, 32, 32, 12],
[nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[12, 32, 32, 12],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(12)]) for _ in range(4)
], 2, 2, 3, p)
bijective(f, decimal=3)
def innerBuilder(num):
maskList = []
for i in range(nlayers):
if i % 2 == 0:
b = torch.zeros(num)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[i] = 1
b = b.reshape(1, num)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
fl = flow.RNVP(maskList, [
utils.SimpleMLPreshape([num] + [hidden] * nmlp + [num],
[nn.Softplus()] * nmlp + [None])
for _ in range(nlayers)
], [
utils.SimpleMLPreshape(
[num] + [hidden] * nmlp + [num],
[nn.Softplus()] * nmlp + [utils.ScalableTanh(num)])
for _ in range(nlayers)
])
return fl
def test_saveload():
p = source.Gaussian([4])
f1 = flow.Scaling(4, [2, 3])
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 4)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
f2 = flow.NICE(maskList, [
utils.SimpleMLPreshape([4, 32, 32, 4],
[nn.ELU(), nn.ELU(), None]) for _ in range(4)
])
f = flow.FlowNet([f1, f2], p)
f1 = flow.Scaling(4)
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 4)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
f2 = flow.NICE(maskList, [
utils.SimpleMLPreshape([4, 32, 32, 4],
[nn.ELU(), nn.ELU(), None]) for _ in range(4)
])
blankf = flow.FlowNet([f1, f2], p)
saveload(f, blankf)
def test_bijective():
p = source.Gaussian([2,2])
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(1,4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:,i] = 1
b=b.reshape(1,2,2)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
f = flow.NICE(maskList, [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)],p)
bijective(f)
def test_symplectic():
p = source.Gaussian([8])
maskList = []
for n in range(4):
if n %2==0:
b = torch.zeros(4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[i] = 1
b=b.reshape(1,4)
else:
b = 1-b
maskList.append(b)
maskList = torch.cat(maskList,0).to(torch.float32)
fl = flow.RNVP(maskList, [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),None]) for _ in range(4)], [utils.SimpleMLPreshape([4,32,32,4],[nn.ELU(),nn.ELU(),utils.ScalableTanh(4)]) for _ in range(4)])
f = flow.PointTransformation(fl,p)
symplectic(f)
def test_integer():
decimal = flow.ScalingNshifting(256, -128)
p = source.MixtureDiscreteLogistic([3, 32, 32], 5, decimal, utils.roundingWidentityGradient)
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.cat([torch.zeros(3 * 32 * 16), torch.ones(3 * 32 * 16)])[torch.randperm(3 * 32 * 32)].reshape(1, 3, 32, 32)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
tList = [utils.SimpleMLPreshape([3 * 32 * 16, 200, 500, 3 * 32 * 16], [nn.ELU(), nn.ELU(), None]) for _ in range(4)]
f = flow.DiscreteNICE(maskList, tList, decimal, utils.roundingWidentityGradient, p)
x, _ = f.sample(100)
assert np.all(np.equal(np.mod(x.detach().numpy(), 1), 0))
zx, _ = f.inverse(x)
assert np.all(np.equal(np.mod(zx.detach().numpy(), 1), 0))
xzx, _ = f.forward(zx)
assert np.all(np.equal(np.mod(xzx.detach().numpy(), 1), 0))
def test_saveload():
p = source.Gaussian([1, 4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
repeat = 2
t = flow.MERA(2, length, layers, repeat, prior=p)
p = source.Gaussian([1, 4, 4])
BigList = []
for _ in range(2 * 2 * 2):
maskList = []
for n in range(4):
if n % 2 == 0:
b = torch.zeros(1, 4)
i = torch.randperm(b.numel()).narrow(0, 0, b.numel() // 2)
b.zero_()[:, i] = 1
b = b.reshape(1, 1, 2, 2)
else:
b = 1 - b
maskList.append(b)
maskList = torch.cat(maskList, 0).to(torch.float32)
BigList.append(maskList)
layers = [
flow.RNVP(BigList[n], [
utils.SimpleMLPreshape(
[4, 32, 32, 4], [nn.ELU(), nn.ELU(), None]) for _ in range(4)
], [
utils.SimpleMLPreshape(
[4, 32, 32, 4],
[nn.ELU(), nn.ELU(), utils.ScalableTanh(4)]) for _ in range(4)
]) for n in range(2 * 2 * 2)
]
length = 4
repeat = 2
blankt = flow.MERA(2, length, layers, repeat, prior=p)
saveload(t, blankt)
