Source code for mmedit.models.backbones.encoder_decoders.necks.contextual_attention_neck

import torch.nn as nn
from mmcv.cnn import ConvModule

from mmedit.models.common import SimpleGatedConvModule
from mmedit.models.common.contextual_attention import ContextualAttentionModule
from mmedit.models.registry import COMPONENTS


[docs]@COMPONENTS.register_module() class ContextualAttentionNeck(nn.Module): """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). """ _conv_type = dict(conv=ConvModule, gated_conv=SimpleGatedConvModule) 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(ContextualAttentionNeck, self).__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