Source code for mmagic.models.editors.deepfillv1.contextual_attention_neck
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.cnn import ConvModule
from mmengine.model import BaseModule
from mmagic.models.archs import SimpleGatedConvModule
from mmagic.models.editors.deepfillv1.contextual_attention import \
ContextualAttentionModule
from mmagic.registry import MODELS
@MODELS.register_module()
[docs]class ContextualAttentionNeck(BaseModule):
"""Neck with contextual attention module.
Args:
in_channels (int): The number of input channels.
conv_type (str): The type of conv module. In DeepFillv1 model, the
`conv_type` should be 'conv'. In DeepFillv2 model, the `conv_type`
should be 'gated_conv'.
conv_cfg (dict | None): Config of conv module. Default: None.
norm_cfg (dict | None): Config of norm module. Default: None.
act_cfg (dict | None): Config of activation layer. Default:
dict(type='ELU').
contextual_attention_args (dict): Config of contextual attention
module. Default: dict(softmax_scale=10.).
kwargs (keyword arguments).
"""
def __init__(self,
in_channels,
conv_type='conv',
conv_cfg=None,
norm_cfg=None,
act_cfg=dict(type='ELU'),
contextual_attention_args=dict(softmax_scale=10.),
**kwargs):
super().__init__()
self.contextual_attention = ContextualAttentionModule(
**contextual_attention_args)
conv_module = self._conv_type[conv_type]
self.conv1 = conv_module(
in_channels,
in_channels,
3,
padding=1,
conv_cfg=conv_cfg,
norm_cfg=norm_cfg,
act_cfg=act_cfg,
**kwargs)
self.conv2 = conv_module(
in_channels,
in_channels,
3,
padding=1,
conv_cfg=conv_cfg,
norm_cfg=norm_cfg,
act_cfg=act_cfg,
**kwargs)
[docs] def forward(self, x, mask):
"""Forward Function.
Args:
x (torch.Tensor): Input tensor with shape of (n, c, h, w).
mask (torch.Tensor): Input tensor with shape of (n, 1, h, w).
Returns:
torch.Tensor: Output tensor with shape of (n, c, h', w').
"""
x, offset = self.contextual_attention(x, x, mask)
x = self.conv1(x)
x = self.conv2(x)
return x, offset