def test_invertible_module_wrapper_simple_inverse(coupling):
"""InvertibleModuleWrapper inverse test"""
for seed in range(10):
set_seeds(seed)
# define some data
X = torch.rand(2, 4, 5, 5).requires_grad_()
# define an arbitrary reversible function
coupling_fn = create_coupling(Fm=torch.nn.Conv2d(2, 2, 3, padding=1),
coupling=coupling,
implementation_fwd=-1,
implementation_bwd=-1,
adapter=AffineAdapterNaive)
fn = InvertibleModuleWrapper(fn=coupling_fn,
keep_input=False,
keep_input_inverse=False)
# compute output
Y = fn.forward(X.clone())
# compute input from output
X2 = fn.inverse(Y)
# check that the inverted output and the original input are approximately similar
assert torch.allclose(X2.detach(), X.detach(), atol=1e-06)
def test_input_output_invertible_function_share_tensor():
fn = IdentityInverse()
rm = InvertibleModuleWrapper(fn=fn, keep_input=True, keep_input_inverse=True)
X = torch.rand(1, 2, 5, 5, dtype=torch.float32).requires_grad_()
assert not is_invertible_module(fn, test_input_shape=X.shape, atol=1e-6)
with pytest.raises(RuntimeError):
rm.forward(X)
fn.multiply_forward = True
rm.forward(X)
assert not is_invertible_module(fn, test_input_shape=X.shape, atol=1e-6)
with pytest.raises(RuntimeError):
rm.inverse(X)
fn.multiply_inverse = True
rm.inverse(X)
assert is_invertible_module(fn, test_input_shape=X.shape, atol=1e-6)
def __init__(self, Gm, coupling='additive', depth=10, implementation_fwd=-1, implementation_bwd=-1,
keep_input=False, adapter=None):
super(SubModuleStack, self).__init__()
fn = create_coupling(Fm=Gm, Gm=Gm, coupling=coupling, implementation_fwd=implementation_fwd, implementation_bwd=implementation_bwd, adapter=adapter)
self.stack = torch.nn.ModuleList(
[InvertibleModuleWrapper(fn=fn, keep_input=keep_input, keep_input_inverse=keep_input) for _ in range(depth)]
)
def test_multi(disable):
split = InvertibleModuleWrapper(SplitChannels(2), disable = disable)
concat = InvertibleModuleWrapper(ConcatenateChannels(2), disable = disable)
assert is_invertible_module(split, test_input_shape=(1, 3, 32, 32))
assert is_invertible_module(concat, test_input_shape=((1, 2, 32, 32), (1, 1, 32, 32)))
conv_a = torch.nn.Conv2d(2, 2, 3)
conv_b = torch.nn.Conv2d(1, 1, 3)
x = torch.rand(1, 3, 32, 32)
x.requires_grad = True
a, b = split(x)
a, b = conv_a(a), conv_b(b)
y = concat(a, b)
loss = torch.sum(y)
loss.backward()
def __init__(self,
inplanes,
planes,
stride=1,
downsample=None,
noactivation=False):
super(RevBasicBlock, self).__init__()
if downsample is None and stride == 1:
gm = BasicBlockSub(inplanes // 2, planes // 2, stride,
noactivation)
fm = BasicBlockSub(inplanes // 2, planes // 2, stride,
noactivation)
coupling = create_coupling(Fm=fm, Gm=gm, coupling='additive')
self.revblock = InvertibleModuleWrapper(fn=coupling,
keep_input=False)
else:
self.basicblock_sub = BasicBlockSub(inplanes, planes, stride,
noactivation)
self.downsample = downsample
self.stride = stride
def test_normal_vs_invertible_module_wrapper(coupling):
"""InvertibleModuleWrapper test if similar gradients and weights results are obtained after similar training"""
for seed in range(10):
set_seeds(seed)
X = torch.rand(2, 4, 5, 5)
# define models and their copies
c1 = torch.nn.Conv2d(2, 2, 3, padding=1)
c2 = torch.nn.Conv2d(2, 2, 3, padding=1)
c1_2 = copy.deepcopy(c1)
c2_2 = copy.deepcopy(c2)
# are weights between models the same, but do they differ between convolutions?
assert torch.equal(c1.weight, c1_2.weight)
assert torch.equal(c2.weight, c2_2.weight)
assert torch.equal(c1.bias, c1_2.bias)
assert torch.equal(c2.bias, c2_2.bias)
assert not torch.equal(c1.weight, c2.weight)
# define optimizers
optim1 = torch.optim.SGD([e for e in c1.parameters()] + [e for e in c2.parameters()], 0.1)
optim2 = torch.optim.SGD([e for e in c1_2.parameters()] + [e for e in c2_2.parameters()], 0.1)
for e in [c1, c2, c1_2, c2_2]:
e.train()
# define an arbitrary reversible function and define graph for model 1
Xin = X.clone().requires_grad_()
coupling_fn = create_coupling(Fm=c1_2, Gm=c2_2, coupling=coupling, implementation_fwd=-1,
implementation_bwd=-1, adapter=AffineAdapterNaive)
fn = InvertibleModuleWrapper(fn=coupling_fn, keep_input=False, keep_input_inverse=False)
Y = fn.forward(Xin)
loss2 = torch.mean(Y)
# define the reversible function without custom backprop and define graph for model 2
XX = X.clone().detach().requires_grad_()
x1, x2 = torch.chunk(XX, 2, dim=1)
if coupling == 'additive':
y1 = x1 + c1.forward(x2)
y2 = x2 + c2.forward(y1)
elif coupling == 'affine':
fmr2 = c1.forward(x2)
fmr1 = torch.exp(fmr2)
y1 = (x1 * fmr1) + fmr2
gmr2 = c2.forward(y1)
gmr1 = torch.exp(gmr2)
y2 = (x2 * gmr1) + gmr2
else:
raise NotImplementedError()
YY = torch.cat([y1, y2], dim=1)
loss = torch.mean(YY)
# compute gradients manually
grads = torch.autograd.grad(loss, (XX, c1.weight, c2.weight, c1.bias, c2.bias), None, retain_graph=True)
# compute gradients and perform optimization model 2
loss.backward()
optim1.step()
# gradients computed manually match those of the .backward() pass
assert torch.equal(c1.weight.grad, grads[1])
assert torch.equal(c2.weight.grad, grads[2])
assert torch.equal(c1.bias.grad, grads[3])
assert torch.equal(c2.bias.grad, grads[4])
# weights differ after training a single model?
assert not torch.equal(c1.weight, c1_2.weight)
assert not torch.equal(c2.weight, c2_2.weight)
assert not torch.equal(c1.bias, c1_2.bias)
assert not torch.equal(c2.bias, c2_2.bias)
# compute gradients and perform optimization model 1
loss2.backward()
optim2.step()
# input is contiguous tests
assert Xin.is_contiguous()
assert Y.is_contiguous()
# weights are approximately the same after training both models?
assert torch.allclose(c1.weight.detach(), c1_2.weight.detach())
assert torch.allclose(c2.weight.detach(), c2_2.weight.detach())
assert torch.allclose(c1.bias.detach(), c1_2.bias.detach())
assert torch.allclose(c2.bias.detach(), c2_2.bias.detach())
# gradients are approximately the same after training both models?
assert torch.allclose(c1.weight.grad.detach(), c1_2.weight.grad.detach())
assert torch.allclose(c2.weight.grad.detach(), c2_2.weight.grad.detach())
assert torch.allclose(c1.bias.grad.detach(), c1_2.bias.grad.detach())
assert torch.allclose(c2.bias.grad.detach(), c2_2.bias.grad.detach())
def test_invertible_module_wrapper_disabled_versus_enabled(inverted):
set_seeds(42)
Gm = SubModule(in_filters=5, out_filters=5)
coupling_fn = create_coupling(Fm=Gm, Gm=Gm, coupling='additive', implementation_fwd=-1,
implementation_bwd=-1)
rb = InvertibleModuleWrapper(fn=coupling_fn, keep_input=False, keep_input_inverse=False)
rb2 = InvertibleModuleWrapper(fn=copy.deepcopy(coupling_fn), keep_input=False, keep_input_inverse=False)
rb.eval()
rb2.eval()
rb2.disable = True
with torch.no_grad():
dims = (2, 10, 8, 8)
data = torch.rand(*dims, dtype=torch.float32)
X, X2 = data.clone().detach().requires_grad_(), data.clone().detach().requires_grad_()
if not inverted:
Y = rb(X)
Y2 = rb2(X2)
else:
Y = rb.inverse(X)
Y2 = rb2.inverse(X2)
assert torch.allclose(Y, Y2)
assert is_memory_cleared(X, True, dims)
assert is_memory_cleared(X2, False, dims)
def test_invertible_module_wrapper_fwd_bwd(fn, bwd, keep_input, keep_input_inverse):
"""InvertibleModuleWrapper tests for the memory saving forward and backward passes
* test inversion Y = RB(X) and X = RB.inverse(Y)
* test training the block for a single step and compare weights for implementations: 0, 1
* test automatic discard of input X and its retrieval after the backward pass
* test usage of BN to identify non-contiguous memory blocks
"""
for seed in range(10):
set_seeds(seed)
dims = (2, 10, 8, 8)
data = torch.rand(*dims, dtype=torch.float32)
target_data = torch.rand(*dims, dtype=torch.float32)
assert is_invertible_module(fn, test_input_shape=data.shape, atol=1e-4)
# test with zero padded convolution
with torch.set_grad_enabled(True):
X = data.clone().requires_grad_()
Ytarget = target_data.clone()
Xshape = X.shape
rb = InvertibleModuleWrapper(fn=fn, keep_input=keep_input, keep_input_inverse=keep_input_inverse)
s_grad = [p.detach().clone() for p in rb.parameters()]
rb.train()
rb.zero_grad()
optim = torch.optim.RMSprop(rb.parameters())
optim.zero_grad()
if not bwd:
Xin = X.clone().requires_grad_()
Y = rb(Xin)
Yrev = Y.detach().clone().requires_grad_()
Xinv = rb.inverse(Yrev)
else:
Xin = X.clone().requires_grad_()
Y = rb.inverse(Xin)
Yrev = Y.detach().clone().requires_grad_()
Xinv = rb(Yrev)
loss = torch.nn.MSELoss()(Y, Ytarget)
# has input been retained/discarded after forward (and backward) passes?
if not bwd:
assert is_memory_cleared(Yrev, not keep_input_inverse, Xshape)
assert is_memory_cleared(Xin, not keep_input, Xshape)
else:
assert is_memory_cleared(Xin, not keep_input_inverse, Xshape)
assert is_memory_cleared(Yrev, not keep_input, Xshape)
optim.zero_grad()
loss.backward()
optim.step()
assert Y.shape == Xshape
assert X.detach().shape == data.shape
assert torch.allclose(X.detach(), data, atol=1e-06)
assert torch.allclose(X.detach(), Xinv.detach(), atol=1e-05) # Model is now trained and will differ
grads = [p.detach().clone() for p in rb.parameters()]
assert not torch.allclose(grads[0], s_grad[0])