parent
14edacdf9a
commit
780860d140
@ -0,0 +1,37 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .evo_norm import EvoNormBatch2d, EvoNormSample2d
|
||||
from .norm_act import BatchNormAct2d
|
||||
try:
|
||||
from inplace_abn import InPlaceABN
|
||||
has_iabn = True
|
||||
except ImportError:
|
||||
has_iabn = False
|
||||
|
||||
|
||||
def create_norm_act(layer_type, num_features, jit=False, **kwargs):
|
||||
layer_parts = layer_type.split('_')
|
||||
assert len(layer_parts) in (1, 2)
|
||||
layer_class = layer_parts[0].lower()
|
||||
#activation_class = layer_parts[1].lower() if len(layer_parts) > 1 else '' # FIXME support string act selection
|
||||
|
||||
if layer_class == "batchnormact":
|
||||
layer = BatchNormAct2d(num_features, **kwargs) # defaults to RELU of no kwargs override
|
||||
elif layer_class == "batchnormrelu":
|
||||
assert 'act_layer' not in kwargs
|
||||
layer = BatchNormAct2d(num_features, act_layer=nn.ReLU, **kwargs)
|
||||
elif layer_class == "evonormbatch":
|
||||
layer = EvoNormBatch2d(num_features, **kwargs)
|
||||
elif layer_class == "evonormsample":
|
||||
layer = EvoNormSample2d(num_features, **kwargs)
|
||||
elif layer_class == "iabn" or layer_class == "inplaceabn":
|
||||
if not has_iabn:
|
||||
raise ImportError(
|
||||
"Pplease install InplaceABN:'pip install git+https://github.com/mapillary/inplace_abn.git@v1.0.11'")
|
||||
layer = InPlaceABN(num_features, **kwargs)
|
||||
else:
|
||||
assert False, "Invalid norm_act layer (%s)" % layer_class
|
||||
if jit:
|
||||
layer = torch.jit.script(layer)
|
||||
return layer
|
Loading…
Reference in new issue