def insert_acblock(model: nn.Module):
items = list(model.named_children())
idx = 0
while idx < len(items):
name, module = items[idx]
if isinstance(module, nn.Conv2d) and module.kernel_size[0] > 1:
# 将标准卷积替换为ACBlock
in_channels = module.in_channels
out_channels = module.out_channels
kernel_size = module.kernel_size
stride = module.stride
padding = module.padding
dilation = module.dilation
groups = module.groups
padding_mode = module.padding_mode
acblock = AsymmetricConvolutionBlock(in_channels,
out_channels,
kernel_size[0],
stride[0],
padding=padding[0],
padding_mode=padding_mode,
dilation=dilation,
groups=groups)
model.add_module(name, acblock)
# 如果conv层之后跟随着BN层，那么删除该BN层
# 参考[About BN layer #35](https://github.com/DingXiaoH/ACNet/issues/35)
if (idx + 1) < len(items) and isinstance(items[idx + 1][1],
nn.BatchNorm2d):
new_layer = nn.Identity()
model.add_module(items[idx + 1][0], new_layer)
else:
insert_acblock(module)
idx += 1
def add_input_quantization(mod: nn.Module):
# Instantiate input observer, for keeping track of the input quantization specs
if mod.qconfig is not None:
mod.add_module('input_quantization', mod.qconfig.activation())
# Only update the observer, but do not apply quantization
mod.input_quantization.enable_fake_quant(False)
mod.register_forward_pre_hook(input_observer_hook)
def convertBatchNorm(cls,
module: nn.Module,
num_micro_batches: int = 1) -> nn.Module:
"""Converts a :class:`nn.BatchNorm` or underlying
:class:`nn.BatchNorm`s into :class:`DelayedBatchNorm`::
from torchvision.models.resnet import resnet101
from pytorch_Gpipe.delayedNorm import DelayedBatchNorm
model = resnet101()
model = DelayedBatchNorm.convertBatchNorm(model)
"""
if isinstance(module, DelayedBatchNorm
) and module.num_micro_batches is num_micro_batches:
return module
if isinstance(module, _BatchNorm) and module.track_running_stats:
module_output = DelayedBatchNorm(module.num_features, module.eps,
module.momentum, module.affine,
num_micro_batches)
if module.affine:
module_output.register_parameter('weight', module.weight)
module_output.register_parameter('bias', module.bias)
module_output.register_buffer('running_mean', module.running_mean)
module_output.register_buffer('running_var', module.running_var)
module_output.register_buffer('num_batches_tracked',
module.num_batches_tracked)
return module_output
for name, child in module.named_children():
module.add_module(name,
cls.convertBatchNorm(child, num_micro_batches))
return module
def add_modules(target_module: nn.Module,
module_list: List[Union[nn.Module, WithModule]],
prefix: str = 'unit_'):
"""
Adds modules to target modules. WithModule type is supported
Args:
target_module: Module to be injected
module_list: List of modules to be added
prefix: Prefix to module names
"""
for i, unit in enumerate(module_list):
name = f'{prefix}{i}'
if isinstance(unit, nn.Module):
target_module.add_module(name, unit)
elif 'module' in dir(unit) and unit.module is not None:
target_module.add_module(name, unit.module)
def replace_layers_in_module(module: nn.Module, mapping_fn: Callable, *args,
**kwargs) -> bool:
"""
Recursively iterate over the children of a module and replace them according to `mapping_fn`.
Returns:
True if a layer has been changed.
"""
changed = False
for name, child in module.named_children():
new_module = mapping_fn(child, *args, **kwargs)
if new_module is not None:
changed = True
module.add_module(name, new_module)
# recursively apply to child
changed |= replace_layers_in_module(child, mapping_fn, *args, **kwargs)
return changed
def _transform_lstm_rnn(module: Module, debug: bool) -> GraphModule:
"""Transforms multi-layer RNN/LSTM to Sequential of single-layer RNN/LSTM.
Converts multi-layer RNN/LSTM to Sequential with single-layer RNN/LSTM.
If dropout probability is nonzero, creates intermediate dropout layers.
Finally, copies training mode.
Args:
module: container module to transform
debug: whether to print debug messages
Returns:
equivalent transformed module
Raises:
NotImplementedError: if initial hidden state is used in forward pass
"""
if debug:
print("\tBegin transformation: LSTM, RNN")
graph: Graph = BackpackTracer().trace(module)
nodes = [
n
for n in graph.nodes
if n.op == "call_module"
and isinstance(module.get_submodule(n.target), (RNN, LSTM))
and module.get_submodule(n.target).num_layers > 1
]
for node in nodes:
if len(node.args) > 1:
raise NotImplementedError(
"For conversion, LSTM/RNN input must not have hidden states."
)
lstm_module_replace = _make_rnn_backpack(module.get_submodule(node.target))
module.add_module(node.target, lstm_module_replace)
graph.lint()
if debug:
print(f"\tRNNs, LSTMs transformed: {len(nodes)}")
return GraphModule(module, graph)
def fuse_repvgg_block(model: nn.Module):
for name, module in model.named_children():
if isinstance(module, RepVGGBlock):
# 将RepVGGBlock替换为标准卷积
kernel, bias = get_equivalent_kernel_bias(
module.rbr_dense, module.rbr_1x1, module.rbr_identity,
module.in_channels, module.groups, module.padding)
# 新建标准卷积，赋值权重和偏差后重新插入模型
fused_conv = nn.Conv2d(module.in_channels,
module.out_channels,
module.kernel_size,
stride=module.stride,
padding=module.padding,
dilation=module.dilation,
groups=module.groups,
padding_mode=module.padding_mode,
bias=True)
fused_conv.weight = nn.Parameter(kernel.detach().cpu())
fused_conv.bias = nn.Parameter(bias.detach().cpu())
model.add_module(name, fused_conv)
else:
fuse_repvgg_block(module)
def modify_for_grad_receptive_field(module: nn.Module):
for name, m in module.named_children():
if mod_isinstance(m, nn.Conv2d):
init_for_grad_receptive_field(m)
elif mod_isinstance(m, nn.MaxPool2d):
# change maxpool to avgpool
replacement = nn.AvgPool2d(
kernel_size=cast(int, m.kernel_size),
stride=cast(int, m.stride),
padding=cast(int, m.padding),
ceil_mode=cast(bool, m.ceil_mode),
)
module.add_module(name, replacement)
elif mod_isinstance(m, nn.Dropout):
# turn off Dropout
# https://discuss.pytorch.org/t/how-to-freeze-bn-layers-while-training-the-rest-of-network-mean-and-var-wont-freeze/89736/10
m.eval()
elif mod_isinstance(m, nn.BatchNorm2d):
# turn off batchnorm
# https://discuss.pytorch.org/t/how-to-close-batchnorm-when-using-torchvision-models/21812/2
if hasattr(m, "reset_parameters"):
cast(Callable, m.reset_parameters)()
m.eval()
init_for_grad_receptive_field(m)
class FCDenseNet(Module):
r"""
The One Hundred Layers Tiramisu: Fully Convolutional DenseNets for Semantic Segmentation
https://arxiv.org/abs/1611.09326
In this paper, we extend DenseNets to deal with the problem of semantic segmentation. We achieve state-of-the-art
results on urban scene benchmark datasets such as CamVid and Gatech, without any further post-processing module nor
pretraining. Moreover, due to smart construction of the model, our approach has much less parameters than currently
published best entries for these datasets.
"""
def __init__(
self,
in_channels: int = 3,
out_channels: int = 1000,
initial_num_features: int = 48,
dropout: float = 0.2,
down_dense_growth_rates: Union[int, Sequence[int]] = 16,
down_dense_bottleneck_ratios: Union[Optional[int],
Sequence[Optional[int]]] = None,
down_dense_num_layers: Union[int, Sequence[int]] = (4, 5, 7, 10, 12),
down_transition_compression_factors: Union[float,
Sequence[float]] = 1.0,
middle_dense_growth_rate: int = 16,
middle_dense_bottleneck: Optional[int] = None,
middle_dense_num_layers: int = 15,
up_dense_growth_rates: Union[int, Sequence[int]] = 16,
up_dense_bottleneck_ratios: Union[Optional[int],
Sequence[Optional[int]]] = None,
up_dense_num_layers: Union[int, Sequence[int]] = (12, 10, 7, 5, 4)):
super(FCDenseNet, self).__init__()
# region Parameters handling
self.in_channels = in_channels
self.out_channels = out_channels
if type(down_dense_growth_rates) == int:
down_dense_growth_rates = (down_dense_growth_rates, ) * 5
if down_dense_bottleneck_ratios is None or type(
down_dense_bottleneck_ratios) == int:
down_dense_bottleneck_ratios = (down_dense_bottleneck_ratios, ) * 5
if type(down_dense_num_layers) == int:
down_dense_num_layers = (down_dense_num_layers, ) * 5
if type(down_transition_compression_factors) == float:
down_transition_compression_factors = (
down_transition_compression_factors, ) * 5
if type(up_dense_growth_rates) == int:
up_dense_growth_rates = (up_dense_growth_rates, ) * 5
if up_dense_bottleneck_ratios is None or type(
up_dense_bottleneck_ratios) == int:
up_dense_bottleneck_ratios = (up_dense_bottleneck_ratios, ) * 5
if type(up_dense_num_layers) == int:
up_dense_num_layers = (up_dense_num_layers, ) * 5
# endregion
# region First convolution
# The Lasagne implementation uses convolution with 'same' padding, the PyTorch equivalent is padding=1
self.features = Conv2d(in_channels,
initial_num_features,
kernel_size=3,
padding=1,
bias=False)
current_channels = self.features.out_channels
# endregion
# region Downward path
# Pairs of Dense Blocks with input concatenation and TransitionDown layers
down_dense_params = [{
'concat_input': True,
'growth_rate': gr,
'num_layers': nl,
'dense_layer_params': {
'dropout': dropout,
'bottleneck_ratio': br
}
} for gr, nl, br in zip(down_dense_growth_rates, down_dense_num_layers,
down_dense_bottleneck_ratios)]
down_transition_params = [{
'dropout': dropout,
'compression': c
} for c in down_transition_compression_factors]
skip_connections_channels = []
self.down_dense = Module()
self.down_trans = Module()
down_pairs_params = zip(down_dense_params, down_transition_params)
for i, (dense_params,
transition_params) in enumerate(down_pairs_params):
block = DenseBlock(current_channels, **dense_params)
current_channels = block.out_channels
self.down_dense.add_module(f'block_{i}', block)
skip_connections_channels.append(block.out_channels)
transition = TransitionDown(current_channels, **transition_params)
current_channels = transition.out_channels
self.down_trans.add_module(f'trans_{i}', transition)
# endregion
# region Middle block
# Renamed from "bottleneck" in the paper, to avoid confusion with the Bottleneck of DenseLayers
self.middle = DenseBlock(current_channels,
middle_dense_growth_rate,
middle_dense_num_layers,
concat_input=True,
dense_layer_params={
'dropout': dropout,
'bottleneck_ratio':
middle_dense_bottleneck
})
current_channels = self.middle.out_channels
# endregion
# region Upward path
# Pairs of TransitionUp layers and Dense Blocks without input concatenation
up_transition_params = [{
'skip_channels': sc,
} for sc in reversed(skip_connections_channels)]
up_dense_params = [{
'concat_input': False,
'growth_rate': gr,
'num_layers': nl,
'dense_layer_params': {
'dropout': dropout,
'bottleneck_ratio': br
}
} for gr, nl, br in zip(up_dense_growth_rates, up_dense_num_layers,
up_dense_bottleneck_ratios)]
self.up_dense = Module()
self.up_trans = Module()
up_pairs_params = zip(up_transition_params, up_dense_params)
for i, (transition_params_up,
dense_params_up) in enumerate(up_pairs_params):
transition = TransitionUp(current_channels, **transition_params_up)
current_channels = transition.out_channels
self.up_trans.add_module(f'trans_{i}', transition)
block = DenseBlock(current_channels, **dense_params_up)
current_channels = block.out_channels
self.up_dense.add_module(f'block_{i}', block)
# endregion
# region Final convolution
self.final = Conv2d(current_channels,
out_channels,
kernel_size=1,
bias=False)
# endregion
# region Weight initialization
for module in self.modules():
if isinstance(module, Conv2d):
init.kaiming_normal_(module.weight)
elif isinstance(module, BatchNorm2d):
module.reset_parameters()
elif isinstance(module, Linear):
init.xavier_uniform(module.weight)
init.constant(module.bias, 0)
# endregion
def forward(self, x):
res = self.features(x)
skip_tensors = []
for dense, trans in zip(self.down_dense.children(),
self.down_trans.children()):
res = dense(res)
skip_tensors.append(res)
res = trans(res)
res = self.middle(res)
for skip, trans, dense in zip(reversed(skip_tensors),
self.up_trans.children(),
self.up_dense.children()):
res = trans(res, skip)
res = dense(res)
res = self.final(res)
return res
def predict(self, x):
logits = self(x)
return F.softmax(logits)
def _unwrap_layers(module: nn.Module):
for name, sub_module in module.named_children():
if isinstance(sub_module, Wrapper):
module.add_module(name, sub_module.layer)
else:
_unwrap_layers(sub_module)
def fuse_acblock(model: nn.Module, eps=1e-5):
for name, module in model.named_children():
if isinstance(module, AsymmetricConvolutionBlock):
# 将ACBlock替换为标准卷积
# 获取Nx1卷积的权重以及对应BN的权重、偏置、运行时均值、运行时方差
vertical_conv_weight = module.ver_conv.weight
vertical_bn_weight = module.ver_bn.weight
vertical_bn_bias = module.ver_bn.bias
vertical_bn_running_mean = module.ver_bn.running_mean
vertical_bn_running_std = torch.sqrt(module.ver_bn.running_var +
eps)
# 获取1xN卷积的权重以及对应BN的权重、偏置、运行时均值、运行时方差
horizontal_conv_weight = module.hor_conv.weight
horizontal_bn_weight = module.hor_bn.weight
horizontal_bn_bias = module.hor_bn.bias
horizontal_bn_running_mean = module.hor_bn.running_mean
horizontal_bn_running_std = torch.sqrt(module.hor_bn.running_var +
eps)
if isinstance(module.square_bn, nn.Identity):
# 获取NxN卷积的权重以及对应BN的权重、偏置、运行时均值、运行时方差
# 在insert_repvgg_block过程中将Conv后面的BN设置nn.Identity
square_weight = module.square_conv.weight
square_bias = module.square_conv.bias
# 计算偏差
fused_bias = square_bias + vertical_bn_bias + horizontal_bn_bias \
- vertical_bn_running_mean * vertical_bn_weight / vertical_bn_running_std \
- horizontal_bn_running_mean * horizontal_bn_weight / horizontal_bn_running_std
# 计算权重
fused_kernel = square_weight
else:
square_conv_weight = module.square_conv.weight
square_bn_weight = module.square_bn.weight
square_bn_bias = module.square_bn.bias
square_bn_running_mean = module.square_bn.running_mean
square_bn_running_std = torch.sqrt(
module.square_bn.running_var + eps)
# 计算偏差
fused_bias = square_bn_bias + vertical_bn_bias + horizontal_bn_bias \
- square_bn_running_mean * square_bn_weight / square_bn_running_std \
- vertical_bn_running_mean * vertical_bn_weight / vertical_bn_running_std \
- horizontal_bn_running_mean * horizontal_bn_weight / horizontal_bn_running_std
# 计算权重
fused_kernel = _fuse_kernel(square_conv_weight,
square_bn_weight,
square_bn_running_std)
# 计算权重
_add_to_square_kernel(
fused_kernel,
_fuse_kernel(vertical_conv_weight, vertical_bn_weight,
vertical_bn_running_std))
_add_to_square_kernel(
fused_kernel,
_fuse_kernel(horizontal_conv_weight, horizontal_bn_weight,
horizontal_bn_running_std))
# 新建标准卷积，赋值权重和偏差后重新插入模型
fused_conv = nn.Conv2d(module.in_channels,
module.out_channels,
module.kernel_size,
stride=module.stride,
padding=module.padding,
dilation=module.dilation,
groups=module.groups,
padding_mode=module.padding_mode)
fused_conv.weight = nn.Parameter(fused_kernel.detach().cpu())
fused_conv.bias = nn.Parameter(fused_bias.detach().cpu())
model.add_module(name, fused_conv)
else:
fuse_acblock(module, eps=eps)
def _make_dense_layers(module: nn.Module, name, input_dim, output_dim):
module.add_module(name, nn.Linear(input_dim, output_dim))
module.add_module(name+'_bn', nn.BatchNorm1d(output_dim))
def build_regularized_relu_block_(module_block: nn.Module,
suffix: str,
reg: Regularization,
dropout: float = 0.5,
num_elem: int = 0):
if reg == Regularization.DROPOUT:
module_block.add_module('Dropout_{}'.format(suffix),
nn.Dropout(p=dropout))
module_block.add_module('ReLU_{}'.format(suffix), nn.ReLU())
elif reg == Regularization.BN_RELU:
module_block.add_module('BN_{}'.format(suffix),
nn.BatchNorm1d(num_elem))
module_block.add_module('ReLU_{}'.format(suffix), nn.ReLU())
elif reg == Regularization.RELU_BN:
module_block.add_module('ReLU_{}'.format(suffix), nn.ReLU())
module_block.add_module('BN_{}'.format(suffix),
nn.BatchNorm1d(num_elem))
