<filename>tests/test_model/test_backbone/test_shufflenetv1_unit.py
# -*- coding: utf-8 -*-

"""
@date: 2020/12/25 上午10:14
@file: test_shufflenetv1_unit.py
@author: zj
@description: 
"""

import torch
import torch.nn as nn
from zcls.model.backbones.shufflenet.shufflenetv1_unit import ShuffleNetV1Unit


def test_shufflenet_v1_unit():
    data = torch.randn(1, 24, 56, 56)
    inplanes = 24
    planes = 384
    groups = 8

    # 下采样，第一个1x1逐点卷积层不进行分组
    stride = 2
    downsample = nn.AvgPool2d(3, stride=stride, padding=1) if stride == 2 else None
    model = ShuffleNetV1Unit(inplanes, planes, groups, stride, downsample, with_group=False)
    print(model)

    outputs = model(data)
    print(outputs.shape)
    assert outputs.shape == (1, planes, 28, 28)

    # 不进行下采样，第一个1x1逐点卷积层进行分组
    stride = 1
    downsample = nn.AvgPool2d(3, stride=stride, padding=1) if stride == 2 else None
    model = ShuffleNetV1Unit(planes, planes, groups, stride, downsample, with_group=True)
    print(model)

    outputs = model(outputs)
    print(outputs.shape)
    assert outputs.shape == (1, planes, 28, 28)


if __name__ == '__main__':
    test_shufflenet_v1_unit()

<filename>tests/test_model/test_backbone/test_shufflenetv1_unit.py
# -*- coding: utf-8 -*-

"""
@date: 2020/12/25 上午10:14
@file: test_shufflenetv1_unit.py
@author: zj
@description: 
"""

import torch
import torch.nn as nn
from zcls.model.backbones.shufflenet.shufflenetv1_unit import ShuffleNetV1Unit


def test_shufflenet_v1_unit():
    data = torch.randn(1, 24, 56, 56)
    inplanes = 24
    planes = 384
    groups = 8

    # 下采样，第一个1x1逐点卷积层不进行分组
    stride = 2
    downsample = nn.AvgPool2d(3, stride=stride, padding=1) if stride == 2 else None
    model = ShuffleNetV1Unit(inplanes, planes, groups, stride, downsample, with_group=False)
    print(model)

    outputs = model(data)
    print(outputs.shape)
    assert outputs.shape == (1, planes, 28, 28)

    # 不进行下采样，第一个1x1逐点卷积层进行分组
    stride = 1
    downsample = nn.AvgPool2d(3, stride=stride, padding=1) if stride == 2 else None
    model = ShuffleNetV1Unit(planes, planes, groups, stride, downsample, with_group=True)
    print(model)

    outputs = model(outputs)
    print(outputs.shape)
    assert outputs.shape == (1, planes, 28, 28)


if __name__ == '__main__':
    test_shufflenet_v1_unit()

# -*- coding: utf-8 -*- @date: 2020/12/25 上午10:14 @file: test_shufflenetv1_unit.py @author: zj @description: # 下采样，第一个1x1逐点卷积层不进行分组 # 不进行下采样，第一个1x1逐点卷积层进行分组