# -*- coding: utf-8 -*-

"""
@date: 2021/5/16 下午10:22
@file: test_shufflenetv1.py
@author: zj
@description: 
"""

import torch

from zcls.config import cfg
from zcls.config.key_word import KEY_OUTPUT
from zcls.model.recognizers.build import build_recognizer


def test_data(model):
    data = torch.randn(1, 3, 224, 224)
    outputs = model(data)[KEY_OUTPUT]
    print(outputs.shape)

    assert outputs.shape == (1, 1000)


def test_shufflenet():
    cfg.merge_from_file('configs/benchmarks/shufflenet/shufflenet_v1_3g2x_zcls_imagenet_224.yaml')
    print(cfg)
    model = build_recognizer(cfg, torch.device('cpu'))
    print(model)

    test_data(model)


if __name__ == '__main__':
    test_shufflenet()

[((18, 11, 18, 38), 'torch.randn', 'torch.randn', ({(18, 23, 18, 24): '1', (18, 26, 18, 27): '3', (18, 29, 18, 32): '224', (18, 34, 18, 37): '224'}, {}), '(1, 3, 224, 224)', False, 'import torch\n'), ((26, 4, 26, 98), 'zcls.config.cfg.merge_from_file', 'cfg.merge_from_file', ({(26, 24, 26, 97): '"""configs/benchmarks/shufflenet/shufflenet_v1_3g2x_zcls_imagenet_224.yaml"""'}, {}), "(\n    'configs/benchmarks/shufflenet/shufflenet_v1_3g2x_zcls_imagenet_224.yaml')", False, 'from zcls.config import cfg\n'), ((28, 34, 28, 53), 'torch.device', 'torch.device', ({(28, 47, 28, 52): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n')]