mmocr/tests/models/textrecog/backbones/test_very_deep_vgg.py
2022-07-21 10:58:04 +08:00

20 lines
434 B
Python

# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import torch
from mmocr.models.textrecog.backbones import VeryDeepVgg
class TestVeryDeepVgg(TestCase):
def test_forward(self):
model = VeryDeepVgg()
model.init_weights()
model.train()
imgs = torch.randn(1, 3, 32, 160)
feats = model(imgs)
self.assertEqual(feats.shape, torch.Size([1, 512, 1, 41]))