Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import torch.nn as nn | |
| from mmcv.runner import BaseModule | |
| from mmocr.models.builder import BACKBONES | |
| class NRTRModalityTransform(BaseModule): | |
| def __init__(self, | |
| input_channels=3, | |
| init_cfg=[ | |
| dict(type='Kaiming', layer='Conv2d'), | |
| dict(type='Uniform', layer='BatchNorm2d') | |
| ]): | |
| super().__init__(init_cfg=init_cfg) | |
| self.conv_1 = nn.Conv2d( | |
| in_channels=input_channels, | |
| out_channels=32, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1) | |
| self.relu_1 = nn.ReLU(True) | |
| self.bn_1 = nn.BatchNorm2d(32) | |
| self.conv_2 = nn.Conv2d( | |
| in_channels=32, | |
| out_channels=64, | |
| kernel_size=3, | |
| stride=2, | |
| padding=1) | |
| self.relu_2 = nn.ReLU(True) | |
| self.bn_2 = nn.BatchNorm2d(64) | |
| self.linear = nn.Linear(512, 512) | |
| def forward(self, x): | |
| x = self.conv_1(x) | |
| x = self.relu_1(x) | |
| x = self.bn_1(x) | |
| x = self.conv_2(x) | |
| x = self.relu_2(x) | |
| x = self.bn_2(x) | |
| n, c, h, w = x.size() | |
| x = x.permute(0, 3, 2, 1).contiguous().view(n, w, h * c) | |
| x = self.linear(x) | |
| x = x.permute(0, 2, 1).contiguous().view(n, -1, 1, w) | |
| return x | |