|
|
|
from functools import partial |
|
|
|
import torch |
|
|
|
TORCH_VERSION = torch.__version__ |
|
|
|
|
|
def is_rocm_pytorch() -> bool: |
|
is_rocm = False |
|
if TORCH_VERSION != 'parrots': |
|
try: |
|
from torch.utils.cpp_extension import ROCM_HOME |
|
is_rocm = True if ((torch.version.hip is not None) and |
|
(ROCM_HOME is not None)) else False |
|
except ImportError: |
|
pass |
|
return is_rocm |
|
|
|
|
|
def _get_cuda_home(): |
|
if TORCH_VERSION == 'parrots': |
|
from parrots.utils.build_extension import CUDA_HOME |
|
else: |
|
if is_rocm_pytorch(): |
|
from torch.utils.cpp_extension import ROCM_HOME |
|
CUDA_HOME = ROCM_HOME |
|
else: |
|
from torch.utils.cpp_extension import CUDA_HOME |
|
return CUDA_HOME |
|
|
|
|
|
def get_build_config(): |
|
if TORCH_VERSION == 'parrots': |
|
from parrots.config import get_build_info |
|
return get_build_info() |
|
else: |
|
return torch.__config__.show() |
|
|
|
|
|
def _get_conv(): |
|
if TORCH_VERSION == 'parrots': |
|
from parrots.nn.modules.conv import _ConvNd, _ConvTransposeMixin |
|
else: |
|
from torch.nn.modules.conv import _ConvNd, _ConvTransposeMixin |
|
return _ConvNd, _ConvTransposeMixin |
|
|
|
|
|
def _get_dataloader(): |
|
if TORCH_VERSION == 'parrots': |
|
from torch.utils.data import DataLoader, PoolDataLoader |
|
else: |
|
from torch.utils.data import DataLoader |
|
PoolDataLoader = DataLoader |
|
return DataLoader, PoolDataLoader |
|
|
|
|
|
def _get_extension(): |
|
if TORCH_VERSION == 'parrots': |
|
from parrots.utils.build_extension import BuildExtension, Extension |
|
CppExtension = partial(Extension, cuda=False) |
|
CUDAExtension = partial(Extension, cuda=True) |
|
else: |
|
from torch.utils.cpp_extension import (BuildExtension, CppExtension, |
|
CUDAExtension) |
|
return BuildExtension, CppExtension, CUDAExtension |
|
|
|
|
|
def _get_pool(): |
|
if TORCH_VERSION == 'parrots': |
|
from parrots.nn.modules.pool import (_AdaptiveAvgPoolNd, |
|
_AdaptiveMaxPoolNd, _AvgPoolNd, |
|
_MaxPoolNd) |
|
else: |
|
from torch.nn.modules.pooling import (_AdaptiveAvgPoolNd, |
|
_AdaptiveMaxPoolNd, _AvgPoolNd, |
|
_MaxPoolNd) |
|
return _AdaptiveAvgPoolNd, _AdaptiveMaxPoolNd, _AvgPoolNd, _MaxPoolNd |
|
|
|
|
|
def _get_norm(): |
|
if TORCH_VERSION == 'parrots': |
|
from parrots.nn.modules.batchnorm import _BatchNorm, _InstanceNorm |
|
SyncBatchNorm_ = torch.nn.SyncBatchNorm2d |
|
else: |
|
from torch.nn.modules.instancenorm import _InstanceNorm |
|
from torch.nn.modules.batchnorm import _BatchNorm |
|
SyncBatchNorm_ = torch.nn.SyncBatchNorm |
|
return _BatchNorm, _InstanceNorm, SyncBatchNorm_ |
|
|
|
|
|
_ConvNd, _ConvTransposeMixin = _get_conv() |
|
DataLoader, PoolDataLoader = _get_dataloader() |
|
BuildExtension, CppExtension, CUDAExtension = _get_extension() |
|
_BatchNorm, _InstanceNorm, SyncBatchNorm_ = _get_norm() |
|
_AdaptiveAvgPoolNd, _AdaptiveMaxPoolNd, _AvgPoolNd, _MaxPoolNd = _get_pool() |
|
|
|
|
|
class SyncBatchNorm(SyncBatchNorm_): |
|
|
|
def _check_input_dim(self, input): |
|
if TORCH_VERSION == 'parrots': |
|
if input.dim() < 2: |
|
raise ValueError( |
|
f'expected at least 2D input (got {input.dim()}D input)') |
|
else: |
|
super()._check_input_dim(input) |
|
|
