def test_diagonal_butterfly(self): batch_size = 10 for in_size, out_size in [(9, 15), (15, 9)]: for nblocks in [1, 2, 3]: for complex in [False, True]: for increasing_stride in [True, False]: for diag_first in [True, False]: dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, in_size, dtype=dtype) b = torch_butterfly.Butterfly(in_size, out_size, bias=False, complex=complex, increasing_stride=increasing_stride, nblocks=nblocks) twiddle_copy = b.twiddle.clone() diagonal = torch.randn(in_size if diag_first else out_size, dtype=dtype) out = b(input * diagonal) if diag_first else b(input) * diagonal for inplace in [True, False]: b_copy = copy.deepcopy(b) # otherwise inplace would modify b bd = torch_butterfly.combine.diagonal_butterfly( b_copy, diagonal, diag_first, inplace) out_bd = bd(input) self.assertTrue(torch.allclose(out_bd, out, self.rtol, self.atol)) if not inplace: self.assertTrue(torch.allclose(b.twiddle, twiddle_copy, self.rtol, self.atol))
def test_transpose_conjugate_multiply(self): n = 16 for complex in [False, True]: for increasing_stride in [True, False]: for nblocks in [1, 2, 3]: b = torch_butterfly.Butterfly(n, n, False, complex, increasing_stride, nblocks=nblocks) dtype = torch.float32 if not complex else torch.complex64 input = torch.eye(n, dtype=dtype) matrix = b(input).t() matrix_t = b.forward(input, transpose=True).t() matrix_conj = b.forward(input, conjugate=True).t() matrix_t_conj = b.forward(input, transpose=True, conjugate=True).t() self.assertTrue( torch.allclose(matrix.t(), matrix_t, self.rtol, self.atol), (complex, increasing_stride, nblocks)) self.assertTrue( torch.allclose(matrix.conj(), matrix_conj, self.rtol, self.atol), (complex, increasing_stride, nblocks)) self.assertTrue( torch.allclose(matrix.t().conj(), matrix_t_conj, self.rtol, self.atol), (complex, increasing_stride, nblocks))
def test_butterfly_kronecker(self): batch_size = 10 n1 = 16 n2 = 32 for complex in [False, True]: for increasing_stride in [True, False]: dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, n2, n1, dtype=dtype) b1 = torch_butterfly.Butterfly(n1, n1, bias=False, complex=complex, increasing_stride=increasing_stride) b2 = torch_butterfly.Butterfly(n2, n2, bias=False, complex=complex, increasing_stride=increasing_stride) b_tp = torch_butterfly.combine.TensorProduct(b1, b2) out_tp = b_tp(input) b = torch_butterfly.combine.butterfly_kronecker(b1, b2) out = b(input.reshape(batch_size, n2 * n1)).reshape(batch_size, n2, n1) self.assertTrue(torch.allclose(out, out_tp, self.rtol, self.atol))
def test_butterfly_product(self): batch_size = 10 n = 16 in_size = 13 out_size = 15 for complex in [False, True]: for inc_stride1, inc_stride2 in itertools.product([True, False], [True, False]): dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, in_size, dtype=dtype) b1 = torch_butterfly.Butterfly(in_size, n, bias=False, complex=complex, increasing_stride=inc_stride1) b2 = torch_butterfly.Butterfly(n, out_size, bias=False, complex=complex, increasing_stride=inc_stride2) out = b2(b1(input)) b = torch_butterfly.combine.butterfly_product(b1, b2) out_prod = b(input) self.assertTrue(torch.allclose(out_prod, out, self.rtol, self.atol))
def test_butterfly_bmm(self): batch_size = 10 matrix_batch = 3 for device in ['cpu' ] + ([] if not torch.cuda.is_available() else ['cuda']): for in_size, out_size in [(7, 15), (15, 7)]: for complex in [False, True]: for increasing_stride in [True, False]: for nblocks in [1, 2, 3]: # Test shape b_bmm = torch_butterfly.ButterflyBmm( in_size, out_size, matrix_batch, True, complex, increasing_stride, nblocks=nblocks).to(device) dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, matrix_batch, in_size, dtype=dtype, device=device) output = b_bmm(input) self.assertTrue( output.shape == (batch_size, matrix_batch, out_size), (output.shape, device, (in_size, out_size), nblocks)) # Check that the result is the same as looping over butterflies output_loop = [] for i in range(matrix_batch): b = torch_butterfly.Butterfly( in_size, out_size, True, complex, increasing_stride, nblocks=nblocks).to(device) with torch.no_grad(): b.twiddle.copy_( b_bmm.twiddle[i * b_bmm.nstacks:(i + 1) * b_bmm.nstacks]) b.bias.copy_(b_bmm.bias[i]) output_loop.append(b(input[:, i])) output_loop = torch.stack(output_loop, dim=1) self.assertTrue( torch.allclose(output, output_loop), ((output - output_loop).abs().max().item(), output.shape, device, (in_size, out_size), complex))
def test_flip_increasing_stride(self): batch_size = 10 for n in [16, 64]: for nblocks in [1, 2, 3]: for complex in [False, True]: for increasing_stride in [True, False]: dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, n, dtype=dtype) b = torch_butterfly.Butterfly(n, n, bias=False, complex=complex, increasing_stride=increasing_stride, nblocks=nblocks) b_new = torch_butterfly.combine.flip_increasing_stride(b) self.assertTrue(b_new[1].increasing_stride == (not b.increasing_stride)) self.assertTrue(torch.allclose(b_new(input), b(input), self.rtol, self.atol))
def test_autograd(self): """Check if autograd works (especially for complex), by trying to match a 4x4 matrix. """ size = 4 niters = 10000 true_model = nn.Linear(size, size, bias=False) x = torch.eye(size) with torch.no_grad(): y = true_model(x) for complex in [False, True]: if complex: model = nn.Sequential( torch_butterfly.complex_utils.Real2Complex(), torch_butterfly.Butterfly(size, size, bias=False, complex=complex), torch_butterfly.complex_utils.Complex2Real(), ) else: model = torch_butterfly.Butterfly(size, size, bias=False, complex=complex) with torch.no_grad(): inital_loss = F.mse_loss(model(x), y) optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) for i in range(niters): out = model(x) loss = F.mse_loss(out, y) optimizer.zero_grad() loss.backward() optimizer.step() # At least loss should decrease # print(inital_loss, loss) self.assertTrue(loss.item() < inital_loss.item())
def test_matrix_to_butterfly_factor(self): num_repeats = 10 for n in [2, 16, 64]: for _ in range(num_repeats): log_n = int(math.ceil(math.log2(n))) for log_k in range(1, log_n + 1): b = torch_butterfly.Butterfly(n, n, bias=False, init='identity') factor = torch.randn(n // 2, 2, 2) b.twiddle[0, 0, log_k - 1].copy_(factor) matrix = b(torch.eye(n)).t() factor_out = matrix_to_butterfly_factor( matrix.detach().numpy(), log_k, pytorch_format=True, check_input=True) self.assertTrue(torch.allclose(factor_out, factor))
def test_subtwiddle(self): batch_size = 10 n = 16 input_size = 8 for complex in [False, True]: for increasing_stride in [True, False]: for nblocks in [1, 2, 3]: b = torch_butterfly.Butterfly(n, n, True, complex, increasing_stride, nblocks=nblocks) dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, input_size, dtype=dtype) output = b(input, subtwiddle=True) self.assertTrue( output.shape == (batch_size, input_size), (output.shape, n, input_size, complex, nblocks))
def test_butterfly(self): batch_size = 10 for device in ['cpu' ] + ([] if not torch.cuda.is_available() else ['cuda']): for in_size, out_size in [(7, 15), (15, 7)]: for complex in [False, True]: for increasing_stride in [True, False]: for init in ['randn', 'ortho', 'identity']: for nblocks in [1, 2, 3]: b = torch_butterfly.Butterfly( in_size, out_size, True, complex, increasing_stride, init, nblocks=nblocks).to(device) dtype = torch.float32 if not complex else torch.complex64 input = torch.randn(batch_size, in_size, dtype=dtype, device=device) output = b(input) self.assertTrue( output.shape == (batch_size, out_size), (output.shape, device, (in_size, out_size), complex, init, nblocks)) if init == 'ortho': twiddle = b.twiddle if not b.complex else view_as_complex( b.twiddle) twiddle_np = twiddle.detach().to( 'cpu').numpy() twiddle_np = twiddle_np.reshape(-1, 2, 2) twiddle_norm = np.linalg.norm(twiddle_np, ord=2, axis=(1, 2)) self.assertTrue( np.allclose(twiddle_norm, 1), (twiddle_norm, device, (in_size, out_size), complex, init))
def test_butterfly_bmm_tensorproduct(self): # Just to show how to do TensorProduct (e.g., Conv2d) with ButterflyBmm batch_size = 10 in_channels = 3 out_channels = 6 n1, n2 = 32, 16 dtype = torch.complex64 input = torch.randn(batch_size, in_channels, n2, n1, dtype=dtype) # Generate out_channels x in_channels butterfly matrices and loop over them b1s = [ torch_butterfly.Butterfly(n1, n1, bias=False, complex=True) for _ in range(out_channels * in_channels) ] b2s = [ torch_butterfly.Butterfly(n2, n2, bias=False, complex=True) for _ in range(out_channels * in_channels) ] b_tp = [ torch_butterfly.combine.TensorProduct(b1, b2) for b1, b2 in zip(b1s, b2s) ] outputs = [] for o in range(out_channels): output = [] for i in range(in_channels): index = o * in_channels + i output.append(b_tp[index](input[:, i])) outputs.append(torch.stack(output, dim=1)) out = torch.stack(outputs, dim=1) assert out.shape == (batch_size, out_channels, in_channels, n2, n1) # Use ButterflyBmm instead b1_bmm = torch_butterfly.ButterflyBmm(n1, n1, matrix_batch=out_channels * in_channels, bias=False, complex=True) with torch.no_grad(): b1_bmm.twiddle.copy_(torch.cat([b1.twiddle for b1 in b1s])) b2_bmm = torch_butterfly.ButterflyBmm(n2, n2, matrix_batch=out_channels * in_channels, bias=False, complex=True) with torch.no_grad(): b2_bmm.twiddle.copy_(torch.cat([b2.twiddle for b2 in b2s])) input_reshaped = input.transpose(1, 2).reshape(batch_size, n2, 1, in_channels, n1) input_expanded = input_reshaped.expand(batch_size, n2, out_channels, in_channels, n1) out_bmm = b1_bmm( input_expanded.reshape(batch_size, n2, out_channels * in_channels, n1)) out_bmm = out_bmm.transpose( 1, 3) # (batch_size, n1, out_channels * in_channels, n2) out_bmm = b2_bmm( out_bmm) # (batch_size, n1, out_channels * in_channels, n2) out_bmm = out_bmm.permute( 0, 2, 3, 1) # (batch_size, out_channels * in_channels, n2, n1) out_bmm = out_bmm.reshape(batch_size, out_channels, in_channels, n2, n1) self.assertTrue(torch.allclose(out_bmm, out))