👕 Standardize

auto_tutorials_source/tutorial_bayesian.py

@@ -51,7 +51,7 @@
 from functools import partial
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 import torch.optim as optim
 from cli_test_helpers import ArgvContext
 

auto_tutorials_source/tutorial_pe_cifar10.py

@@ -125,7 +125,7 @@ def imshow(img):
 # First we define a vanilla classifier for CIFAR10 for reference. We will use a
 # convolutional neural network.
 
-import torch.nn as nn
+from torch import nn
 import torch.nn.functional as F
 
 

docs/source/quickstart.rst

@@ -48,7 +48,7 @@ trains any ResNet architecture on CIFAR10:
 
     from pathlib import Path
 
-    import torch.nn as nn
+    from torch import nn
 
     from torch_uncertainty import cli_main, init_args
     from torch_uncertainty.baselines import ResNet

experiments/classification/cifar10/resnet.py

@@ -1,7 +1,7 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import ResNet

experiments/classification/cifar10/vgg.py

@@ -1,7 +1,7 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import VGG

experiments/classification/cifar10/wideresnet.py

@@ -1,7 +1,7 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import WideResNet

experiments/classification/cifar100/resnet.py

@@ -1,7 +1,7 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import ResNet

experiments/classification/cifar100/vgg.py

@@ -1,7 +1,7 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import VGG

experiments/classification/cifar100/wideresnet.py

@@ -1,7 +1,7 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import WideResNet

experiments/classification/mnist/bayesian_lenet.py

@@ -2,8 +2,8 @@
 from functools import partial
 from pathlib import Path
 
-import torch.nn as nn
 import torch.optim as optim
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.datamodules import MNISTDataModule

experiments/classification/mnist/lenet.py

@@ -1,8 +1,8 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
 import torch.optim as optim
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.datamodules import MNISTDataModule

experiments/classification/tiny-imagenet/resnet.py

@@ -1,8 +1,8 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
 import torch.optim as optim
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import ResNet

experiments/regression/uci_datasets/mlp-kin8nm.py

@@ -1,8 +1,8 @@
 # fmt: off
 from pathlib import Path
 
-import torch.nn as nn
 import torch.optim as optim
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines.regression.mlp import MLP

tests/_dummies/baseline.py

@@ -2,8 +2,8 @@
 from argparse import ArgumentParser
 from typing import Any
 
-import torch.nn as nn
 from pytorch_lightning import LightningModule
+from torch import nn
 
 from torch_uncertainty.routines.classification import (
     ClassificationEnsemble,

tests/_dummies/model.py

@@ -1,8 +1,7 @@
 # fmt: off
 
 import torch
-import torch.nn as nn
-from torch import Tensor
+from torch import Tensor, nn
 
 # fmt: on
 __all__ = [

tests/baselines/test_batched.py

@@ -1,6 +1,6 @@
 # fmt:off
 import torch
-import torch.nn as nn
+from torch import nn
 from torchinfo import summary
 
 from torch_uncertainty.baselines import ResNet, WideResNet

tests/baselines/test_masked.py

@@ -1,7 +1,7 @@
 # fmt:off
 import pytest
 import torch
-import torch.nn as nn
+from torch import nn
 from torchinfo import summary
 
 from torch_uncertainty.baselines import ResNet, WideResNet

tests/baselines/test_mimo.py

@@ -1,6 +1,6 @@
 # fmt: off
 import torch
-import torch.nn as nn
+from torch import nn
 from torchinfo import summary
 
 from torch_uncertainty.baselines import ResNet, WideResNet

tests/baselines/test_packed.py

@@ -1,7 +1,7 @@
 # fmt:off
 import pytest
 import torch
-import torch.nn as nn
+from torch import nn
 from torchinfo import summary
 
 from torch_uncertainty.baselines import VGG, ResNet, WideResNet

tests/baselines/test_resnet.py

@@ -1,6 +1,6 @@
 # fmt:off
 import pytest
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty.baselines import ResNet, WideResNet
 from torch_uncertainty.optimization_procedures import (

tests/baselines/test_standard.py

@@ -1,6 +1,6 @@
 # fmt:off
 import torch
-import torch.nn as nn
+from torch import nn
 from torchinfo import summary
 
 from torch_uncertainty.baselines import VGG, ResNet, WideResNet

tests/routines/test_classification.py

@@ -2,8 +2,8 @@
 from pathlib import Path
 
 import pytest
-import torch.nn as nn
 from cli_test_helpers import ArgvContext
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.optimization_procedures import optim_cifar10_resnet18

tests/routines/test_regression.py

@@ -1,8 +1,8 @@
 # fmt:off
 from pathlib import Path
 
-import torch.nn as nn
 from cli_test_helpers import ArgvContext
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.optimization_procedures import optim_cifar10_resnet18

tests/test_cli.py

@@ -3,8 +3,8 @@
 from pathlib import Path
 
 import pytest
-import torch.nn as nn
 from cli_test_helpers import ArgvContext
+from torch import nn
 
 from torch_uncertainty import cli_main, init_args
 from torch_uncertainty.baselines import VGG, ResNet, WideResNet

tests/test_losses.py

@@ -1,7 +1,7 @@
 # fmt: off
 import pytest
 import torch
-import torch.nn as nn
+from torch import nn
 
 from torch_uncertainty.layers.bayesian_layers import BayesLinear
 from torch_uncertainty.losses import ELBOLoss

torch_uncertainty/baselines/classification/resnet.py

@@ -4,12 +4,12 @@
 from typing import Any, Literal, Optional, Type, Union
 
 import torch
-import torch.nn as nn
 from pytorch_lightning import LightningModule
 from pytorch_lightning.core.saving import (
     load_hparams_from_tags_csv,
     load_hparams_from_yaml,
 )
+from torch import nn
 
 from ...models.resnet import (
     batched_resnet18,

torch_uncertainty/baselines/classification/vgg.py

@@ -4,12 +4,12 @@
 from typing import Any, Literal, Optional, Type, Union
 
 import torch
-import torch.nn as nn
 from pytorch_lightning import LightningModule
 from pytorch_lightning.core.saving import (
     load_hparams_from_tags_csv,
     load_hparams_from_yaml,
 )
+from torch import nn
 
 from ...models.vgg import (
     packed_vgg11,

torch_uncertainty/baselines/classification/wideresnet.py

@@ -4,12 +4,12 @@
 from typing import Any, Literal, Optional, Type, Union
 
 import torch
-import torch.nn as nn
 from pytorch_lightning import LightningModule
 from pytorch_lightning.core.saving import (
     load_hparams_from_tags_csv,
     load_hparams_from_yaml,
 )
+from torch import nn
 
 from ...models.wideresnet import (
     batched_wideresnet28x10,

torch_uncertainty/baselines/regression/mlp.py

@@ -4,12 +4,12 @@
 from typing import Any, List, Literal, Optional, Union
 
 import torch
-import torch.nn as nn
 from pytorch_lightning import LightningModule
 from pytorch_lightning.core.saving import (
     load_hparams_from_tags_csv,
     load_hparams_from_yaml,
 )
+from torch import nn
 
 from ...models.mlp import mlp, packed_mlp
 from ...routines.regression import RegressionEnsemble, RegressionSingle

torch_uncertainty/datamodules/cifar10.py

@@ -3,10 +3,10 @@
 from pathlib import Path
 from typing import Any, List, Literal, Optional, Union
 
-import torch.nn as nn
 import torchvision.transforms as T
 from pytorch_lightning import LightningDataModule
 from timm.data.auto_augment import rand_augment_transform
+from torch import nn
 from torch.utils.data import DataLoader, Dataset, random_split
 from torchvision.datasets import CIFAR10, SVHN
 

torch_uncertainty/datamodules/cifar100.py

@@ -4,10 +4,10 @@
 from typing import Any, List, Literal, Optional, Union
 
 import torch
-import torch.nn as nn
 import torchvision.transforms as T
 from pytorch_lightning import LightningDataModule
 from timm.data.auto_augment import rand_augment_transform
+from torch import nn
 from torch.utils.data import DataLoader, Dataset, random_split
 from torchvision.datasets import CIFAR100, SVHN
 

torch_uncertainty/datamodules/imagenet.py

@@ -3,11 +3,11 @@
 from pathlib import Path
 from typing import Any, List, Optional, Union
 
-import torch.nn as nn
 import torchvision.transforms as T
 from pytorch_lightning import LightningDataModule
 from timm.data.auto_augment import rand_augment_transform
 from timm.data.mixup import Mixup
+from torch import nn
 from torch.utils.data import DataLoader, Dataset
 from torchvision.datasets import DTD, SVHN, ImageNet, INaturalist
 

torch_uncertainty/datamodules/tiny_imagenet.py

@@ -3,10 +3,10 @@
 from pathlib import Path
 from typing import Any, List, Optional, Union
 
-import torch.nn as nn
 import torchvision.transforms as T
 from pytorch_lightning import LightningDataModule
 from timm.data.auto_augment import rand_augment_transform
+from torch import nn
 from torch.utils.data import DataLoader, Dataset
 from torchvision.datasets import SVHN
 

torch_uncertainty/layers/batchens_layers.py

@@ -3,7 +3,7 @@
 from typing import Optional, Tuple, Union
 
 import torch
-import torch.nn as nn
+from torch import nn
 from torch.nn.common_types import _size_2_t
 from torch.nn.modules.utils import _pair
 

torch_uncertainty/layers/bayesian_layers/sampler.py

@@ -2,8 +2,7 @@
 from typing import Optional
 
 import torch
-import torch.nn as nn
-from torch import Tensor
+from torch import Tensor, nn
 
 import numpy as np
 

torch_uncertainty/layers/masksembles_layers.py

@@ -3,8 +3,7 @@
 from typing import Any, Union
 
 import torch
-import torch.nn as nn
-from torch import Tensor
+from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
 import numpy as np

torch_uncertainty/layers/packed_layers.py

@@ -1,9 +1,8 @@
 # fmt: off
 from typing import Any, Union
 
-import torch.nn as nn
 from einops import rearrange
-from torch import Tensor
+from torch import Tensor, nn
 from torch.nn.common_types import _size_2_t
 
 

torch_uncertainty/models/deep_ensembles.py

@@ -3,7 +3,7 @@
 from typing import List, Optional, Union
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 
 # fmt: on

torch_uncertainty/models/lenet.py

@@ -2,8 +2,8 @@
 from typing import Callable, Dict, Optional, Type, Union
 
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn
 
 from ..layers.bayesian_layers import BayesConv2d, BayesLinear
 from ..layers.packed_layers import PackedConv2d, PackedLinear