mirror of https://github.com/open-mmlab/mmocr.git
64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import torch
|
|
|
|
from mmocr.models.common import (PositionalEncoding, TFDecoderLayer,
|
|
TFEncoderLayer)
|
|
from mmocr.models.textrecog.layers import BasicBlock, Bottleneck
|
|
from mmocr.models.textrecog.layers.conv_layer import conv3x3
|
|
|
|
|
|
def test_conv_layer():
|
|
conv3by3 = conv3x3(3, 6)
|
|
assert conv3by3.in_channels == 3
|
|
assert conv3by3.out_channels == 6
|
|
assert conv3by3.kernel_size == (3, 3)
|
|
|
|
x = torch.rand(1, 64, 224, 224)
|
|
# test basic block
|
|
basic_block = BasicBlock(64, 64)
|
|
assert basic_block.expansion == 1
|
|
|
|
out = basic_block(x)
|
|
|
|
assert out.shape == torch.Size([1, 64, 224, 224])
|
|
|
|
# test bottle neck
|
|
bottle_neck = Bottleneck(64, 64, downsample=True)
|
|
assert bottle_neck.expansion == 4
|
|
|
|
out = bottle_neck(x)
|
|
|
|
assert out.shape == torch.Size([1, 256, 224, 224])
|
|
|
|
|
|
def test_transformer_layer():
|
|
# test decoder_layer
|
|
decoder_layer = TFDecoderLayer()
|
|
in_dec = torch.rand(1, 30, 512)
|
|
out_enc = torch.rand(1, 128, 512)
|
|
out_dec = decoder_layer(in_dec, out_enc)
|
|
assert out_dec.shape == torch.Size([1, 30, 512])
|
|
|
|
decoder_layer = TFDecoderLayer(
|
|
operation_order=('self_attn', 'norm', 'enc_dec_attn', 'norm', 'ffn',
|
|
'norm'))
|
|
out_dec = decoder_layer(in_dec, out_enc)
|
|
assert out_dec.shape == torch.Size([1, 30, 512])
|
|
|
|
# test positional_encoding
|
|
pos_encoder = PositionalEncoding()
|
|
x = torch.rand(1, 30, 512)
|
|
out = pos_encoder(x)
|
|
assert out.size() == x.size()
|
|
|
|
# test encoder_layer
|
|
encoder_layer = TFEncoderLayer()
|
|
in_enc = torch.rand(1, 20, 512)
|
|
out_enc = encoder_layer(in_enc)
|
|
assert out_dec.shape == torch.Size([1, 30, 512])
|
|
|
|
encoder_layer = TFEncoderLayer(
|
|
operation_order=('self_attn', 'norm', 'ffn', 'norm'))
|
|
out_enc = encoder_layer(in_enc)
|
|
assert out_dec.shape == torch.Size([1, 30, 512])
|