Import torch.nn consistently

examples/amoebanetd_speed_benchmark/main.py

@@ -6,8 +6,8 @@
 
 import click
 import torch
-import torch.nn as nn
-from torch.nn import functional as F
+from torch import nn
+import torch.nn.functional as F
 from torch.optim import SGD
 from torch.utils.data import DataLoader
 

examples/resnet/__init__.py

@@ -7,7 +7,7 @@
 from collections import OrderedDict
 from typing import Any, List
 
-import torch.nn as nn
+from torch import nn
 
 from resnet.bottleneck import bottleneck
 from resnet.flatten_sequential import flatten_sequential

examples/resnet/bottleneck.py

@@ -2,8 +2,7 @@
 from collections import OrderedDict
 from typing import TYPE_CHECKING, Optional, Tuple, Union
 
-from torch import Tensor
-import torch.nn as nn
+from torch import Tensor, nn
 
 __all__ = ['bottleneck']
 

examples/resnet101_accuracy_benchmark/main.py

@@ -5,8 +5,8 @@
 
 import click
 import torch
-import torch.nn as nn
-from torch.nn import functional as F
+from torch import nn
+import torch.nn.functional as F
 from torch.optim import SGD
 from torch.optim.lr_scheduler import LambdaLR
 from torch.utils.data import DataLoader
@@ -34,7 +34,7 @@ def dataparallel256(model: nn.Module, devices: List[int]) -> Stuffs:
 
         devices = [devices[0], devices[1]]
         model.to(devices[0])
-        model = torch.nn.DataParallel(model, device_ids=devices, output_device=devices[-1])
+        model = nn.DataParallel(model, device_ids=devices, output_device=devices[-1])
 
         return model, batch_size, [torch.device(device) for device in devices]
 
@@ -44,7 +44,7 @@ def dataparallel1k(model: nn.Module, devices: List[int]) -> Stuffs:
 
         devices = [devices[0], devices[1], devices[2], devices[3]]
         model.to(devices[0])
-        model = torch.nn.DataParallel(model, device_ids=devices, output_device=devices[-1])
+        model = nn.DataParallel(model, device_ids=devices, output_device=devices[-1])
 
         return model, batch_size, [torch.device(device) for device in devices]
 

examples/resnet101_speed_benchmark/main.py

@@ -5,8 +5,8 @@
 
 import click
 import torch
-import torch.nn as nn
-from torch.nn import functional as F
+from torch import nn
+import torch.nn.functional as F
 from torch.optim import SGD
 
 from resnet import resnet101

tests/test_bugs.py

@@ -1,6 +1,6 @@
 import pytest
 import torch
-import torch.nn as nn
+from torch import nn
 
 from torchgpipe import GPipe
 
@@ -25,7 +25,7 @@ def forward(ctx, input):
         def backward(ctx, grad):
             return grad
 
-    class M(torch.nn.Module):
+    class M(nn.Module):
         def forward(self, input):
             return Identity.apply(input)
 

tests/test_deferred_batch_norm.py

@@ -3,8 +3,7 @@
 
 import pytest
 import torch
-import torch.nn as nn
-import torch.optim as optim
+from torch import nn, optim
 
 from torchgpipe.batchnorm import DeferredBatchNorm
 

tests/test_gpipe.py

@@ -4,7 +4,7 @@
 
 import pytest
 import torch
-import torch.nn as nn
+from torch import nn
 
 from torchgpipe import GPipe
 from torchgpipe.gpipe import verify_module

tests/test_inplace.py

@@ -1,6 +1,6 @@
 import pytest
 import torch
-import torch.nn as nn
+from torch import nn
 
 from torchgpipe import GPipe
 

tests/test_pipeline.py

@@ -1,7 +1,7 @@
 import time
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 from torchgpipe.microbatch import Batch
 from torchgpipe.pipeline import Pipeline

tests/test_transparency.py

@@ -1,5 +1,5 @@
 import torch
-import torch.nn as nn
+from torch import nn
 
 from torchgpipe import GPipe
 

torchgpipe/batchnorm.py

@@ -2,8 +2,7 @@
 from typing import Optional, TypeVar, cast
 
 import torch
-from torch import Tensor
-import torch.nn as nn
+from torch import Tensor, nn
 import torch.nn.functional as F
 from torch.nn.modules.batchnorm import _BatchNorm
 

torchgpipe/gpipe.py

@@ -3,10 +3,9 @@
 from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Tuple, Union, cast
 
 import torch
-from torch import Tensor
+from torch import Tensor, nn
 import torch.autograd
 import torch.cuda
-import torch.nn as nn
 
 from torchgpipe.batchnorm import DeferredBatchNorm
 from torchgpipe.microbatch import check, gather, scatter