Update DCT Functions and Fix idct-2/dct-2 Implementations

dct_transform.py

@@ -1,6 +1,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+import pdb
 
 '''adapted from https://github.com/zh217/torch-dct'''
 
@@ -11,18 +12,24 @@ def dct_1d(x):
     :param x: the input signal
     :return: the DCT-II of the signal over the last dimension
     """
+
     x_shape = x.shape
     N = x_shape[-1]
+
+    # reshape the input signal to a 2D tensor with the last dim flattened
     x = x.contiguous().view(-1, N)
 
-    v = torch.cat([x[:, ::2], x[:, 1::2].flip([1])], dim=1)
+    # rearrange cols
+    v = torch.cat([x[:, ::2], x[:, 1::2].flip([1])], dim=-1)
 
-    Vc = torch.rfft(v, 1, onesided=False)
+    # Apply 1-D real-to-complex Fast Fourier Transform along last dim
+    Vc = torch.view_as_real( torch.fft.fft(v, dim=1))
 
     k = - torch.arange(N, dtype=x.dtype, device=x.device)[None, :] * np.pi / (2 * N)
     W_r = torch.cos(k)
     W_i = torch.sin(k)
 
+    # Apply DCT formula
     V = Vc[:, :, 0] * W_r - Vc[:, :, 1] * W_i
     V = 2 * V.view(*x_shape)
 
@@ -54,7 +61,7 @@ def idct_1d(X):
 
     V = torch.cat([V_r.unsqueeze(2), V_i.unsqueeze(2)], dim=2)
 
-    v = torch.irfft(V, 1, onesided=False)
+    v = torch.fft.irfft(torch.view_as_complex(V),n=V.shape[1], dim=1)
     x = v.new_zeros(v.shape)
     x[:, ::2] += v[:, :N - (N // 2)]
     x[:, 1::2] += v.flip([1])[:, :N // 2]
@@ -124,5 +131,3 @@ def output_dims(self, input_dims):
 
         err = torch.abs(x - x_inv).max()
         print(N, err.item(), flush=True)
-
-

inn_architecture.py

@@ -5,6 +5,7 @@
 import torch.nn as nn
 from torch.nn.functional import conv2d, interpolate
 import numpy as np
+import pdb
 
 import FrEIA.framework as Ff
 import FrEIA.modules as Fm
@@ -123,6 +124,7 @@ def fc_constr(c_in, c_out):
     channels = classifier.input_channels
 
     if classifier.dataset == 'MNIST':
+        # import pdb; pdb.set_trace()
         nodes.append(Ff.Node(nodes[-1].out0, Fm.Reshape, {'target_dim':(1, *classifier.dims)}))
         nodes.append(Ff.Node(nodes[-1].out0, Fm.HaarDownsampling, {'rebalance':1.}))
         channels *= 4

train.py

@@ -97,6 +97,8 @@ def log_write(line, endline='\n'):
 
             for i_batch, (x,l) in enumerate(dataset.train_loader):
 
+                import pdb; pdb.set_trace()
+
                 x, y = x.cuda(), dataset.onehot(l.cuda(), label_smoothing)
                 losses = inn(x, y)