Merge branch 'refactor' into dataset_tasks

disent/frameworks/helper/latent_distributions.py

@@ -34,7 +34,7 @@
 
 from disent.frameworks.helper.util import compute_ave_loss
 from disent.nn.loss.kl import kl_loss
-from disent.nn.reductions import loss_reduction
+from disent.nn.loss.reduction import loss_reduction
 
 
 # ========================================================================= #

disent/frameworks/helper/reconstructions.py

@@ -34,8 +34,8 @@
 
 from disent.frameworks.helper.util import compute_ave_loss
 from disent.nn.modules import DisentModule
-from disent.nn.reductions import batch_loss_reduction
-from disent.nn.reductions import loss_reduction
+from disent.nn.loss.reduction import batch_loss_reduction
+from disent.nn.loss.reduction import loss_reduction
 from disent.nn.transform import FftKernel
 
 

disent/frameworks/vae/_unsupervised__dfcvae.py

@@ -36,7 +36,7 @@
 from torchvision.models import vgg19_bn
 from torch.nn import functional as F
 
-from disent.nn.reductions import get_mean_loss_scale
+from disent.nn.loss.reduction import get_mean_loss_scale
 from disent.frameworks.helper.util import compute_ave_loss
 from disent.frameworks.vae._unsupervised__betavae import BetaVae
 from disent.nn.transform.functional import check_tensor

disent/frameworks/vae/_unsupervised__dipvae.py

@@ -30,7 +30,7 @@
 
 from disent.frameworks.helper.util import compute_ave_loss_and_logs
 from disent.frameworks.vae._unsupervised__betavae import BetaVae
-from disent.util.math import torch_cov_matrix
+from disent.nn.functional import torch_cov_matrix
 
 
 # ========================================================================= #

disent/frameworks/vae/experimental/_unsupervised__dorvae.py

@@ -33,10 +33,9 @@
 from disent.frameworks.helper.reconstructions import make_reconstruction_loss
 from disent.frameworks.helper.reconstructions import ReconLossHandler
 from disent.frameworks.vae._supervised__tvae import TripletVae
-from disent.frameworks.vae._unsupervised__betavae import BetaVae
 from disent.frameworks.vae._weaklysupervised__adavae import AdaVae
-from disent.util.math_loss import torch_mse_rank_loss
-from disent.util.math_loss import spearman_rank_loss
+from disent.nn.loss.softsort import torch_mse_rank_loss
+from disent.nn.loss.softsort import spearman_rank_loss
 from experiment.util.hydra_utils import instantiate_recursive
 
 

disent/metrics/_flatness_components.py

@@ -35,8 +35,8 @@
 from disent.metrics._flatness import filter_inactive_factors
 from disent.util import iter_chunks
 from disent.util import to_numpy
-from disent.util.math import torch_mean_generalized
-from disent.util.math import torch_pca
+from disent.nn.functional import torch_mean_generalized
+from disent.nn.functional import torch_pca
 
 
 log = logging.getLogger(__name__)

disent/util/math.py → disent/nn/functional/__init__.py

@@ -21,19 +21,20 @@
 #  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #  SOFTWARE.
 #  ~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~
+
+import logging
 import warnings
 from typing import List
 from typing import Optional
 from typing import Union
 
-import logging
 import numpy as np
 import torch
 
-from disent.util.math_generic import generic_as_int32
-from disent.util.math_generic import generic_max
-from disent.util.math_generic import TypeGenericTensor
-from disent.util.math_generic import TypeGenericTorch
+from disent.nn.functional._generic_tensors import generic_as_int32
+from disent.nn.functional._generic_tensors import generic_max
+from disent.nn.functional._generic_tensors import TypeGenericTensor
+from disent.nn.functional._generic_tensors import TypeGenericTorch
 
 
 log = logging.getLogger(__name__)

disent/nn/transform/_augment.py

@@ -34,9 +34,9 @@
 
 import disent
 from disent.nn.modules import DisentModule
-from disent.util.math import torch_box_kernel_2d
-from disent.util.math import torch_conv2d_channel_wise_fft
-from disent.util.math import torch_gaussian_kernel_2d
+from disent.nn.functional import torch_box_kernel_2d
+from disent.nn.functional import torch_conv2d_channel_wise_fft
+from disent.nn.functional import torch_gaussian_kernel_2d
 
 
 # ========================================================================= #

experiment/exp/05_adversarial_data/run_01_sort_loss.py

@@ -27,8 +27,8 @@
 from torch.utils.data import DataLoader
 
 import experiment.exp.util as H
-from disent.util.math_loss import multi_spearman_rank_loss
-from disent.util.math_loss import torch_soft_rank
+from disent.nn.loss.softsort import multi_spearman_rank_loss
+from disent.nn.loss.softsort import torch_soft_rank
 
 
 # ========================================================================= #

experiment/exp/05_adversarial_data/run_02_check_aug_gt_dists.py

@@ -30,9 +30,9 @@
 
 import torch.nn.functional as F
 import experiment.exp.util as H
-from disent.util.math import torch_conv2d_channel_wise_fft
-from disent.util.math import torch_box_kernel_2d
-from disent.util.math import torch_gaussian_kernel_2d
+from disent.nn.functional import torch_conv2d_channel_wise_fft
+from disent.nn.functional import torch_box_kernel_2d
+from disent.nn.functional import torch_gaussian_kernel_2d
 
 
 # ========================================================================= #

experiment/exp/05_adversarial_data/run_03_train_disentangle_kernel.py

@@ -42,8 +42,8 @@
 from disent.nn.modules import DisentModule
 from disent.util import make_box_str
 from disent.util import seed
-from disent.util.math import torch_conv2d_channel_wise_fft
-from disent.util.math_loss import spearman_rank_loss
+from disent.nn.functional import torch_conv2d_channel_wise_fft
+from disent.nn.loss.softsort import spearman_rank_loss
 from experiment.run import hydra_append_progress_callback
 from experiment.run import hydra_check_cuda
 from experiment.run import hydra_make_logger

experiment/exp/util/_loss.py

@@ -26,7 +26,7 @@
 import torch_optimizer
 from torch.nn import functional as F
 
-from disent.nn.reductions import batch_loss_reduction
+from disent.nn.loss.reduction import batch_loss_reduction
 
 
 # ========================================================================= #

tests/test_math.py

@@ -30,18 +30,18 @@
 
 from disent.data.groundtruth import XYSquaresData
 from disent.dataset.groundtruth import GroundTruthDataset
+from disent.nn.functional import torch_conv2d_channel_wise
+from disent.nn.functional import torch_conv2d_channel_wise_fft
+from disent.nn.functional import torch_corr_matrix
+from disent.nn.functional import torch_cov_matrix
+from disent.nn.functional import torch_dct
+from disent.nn.functional import torch_dct2
+from disent.nn.functional import torch_gaussian_kernel_2d
+from disent.nn.functional import torch_idct
+from disent.nn.functional import torch_idct2
+from disent.nn.functional import torch_mean_generalized
 from disent.nn.transform import ToStandardisedTensor
-from disent.util.math import torch_conv2d_channel_wise
-from disent.util.math import torch_conv2d_channel_wise_fft
 from disent.util import to_numpy
-from disent.util.math import torch_dct
-from disent.util.math import torch_dct2
-from disent.util.math import torch_gaussian_kernel_2d
-from disent.util.math import torch_idct
-from disent.util.math import torch_idct2
-from disent.util.math import torch_corr_matrix
-from disent.util.math import torch_cov_matrix
-from disent.util.math import torch_mean_generalized
 
 
 # ========================================================================= #

tests/test_math_generic.py

@@ -26,11 +26,11 @@
 import pytest
 import torch
 
-from disent.util.math_generic import generic_as_int32
-from disent.util.math_generic import generic_max
-from disent.util.math_generic import generic_min
-from disent.util.math_generic import generic_ndim
-from disent.util.math_generic import generic_shape
+from disent.nn.functional._generic_tensors import generic_as_int32
+from disent.nn.functional._generic_tensors import generic_max
+from disent.nn.functional._generic_tensors import generic_min
+from disent.nn.functional._generic_tensors import generic_ndim
+from disent.nn.functional._generic_tensors import generic_shape
 
 
 # ========================================================================= #

tests/test_transform.py

@@ -26,8 +26,8 @@
 
 from disent.nn.transform import FftGaussianBlur
 from disent.nn.transform._augment import _expand_to_min_max_tuples
-from disent.util.math import torch_gaussian_kernel
-from disent.util.math import torch_gaussian_kernel_2d
+from disent.nn.functional import torch_gaussian_kernel
+from disent.nn.functional import torch_gaussian_kernel_2d
 
 
 # ========================================================================= #