File size: 1,451 Bytes
2784407 07fa01d 2784407 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import sys
from pathlib import Path
import torch
project_root = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(project_root))
from yolo.model.module import SPPELAN, ADown, CBLinear, Conv, Pool
STRIDE = 2
KERNEL_SIZE = 3
IN_CHANNELS = 64
OUT_CHANNELS = 128
NECK_CHANNELS = 64
def test_conv():
conv = Conv(IN_CHANNELS, OUT_CHANNELS, KERNEL_SIZE)
x = torch.randn(1, IN_CHANNELS, 64, 64)
out = conv(x)
assert out.shape == (1, OUT_CHANNELS, 64, 64)
def test_pool_max():
pool = Pool("max", 2, stride=2)
x = torch.randn(1, IN_CHANNELS, 64, 64)
out = pool(x)
assert out.shape == (1, IN_CHANNELS, 32, 32)
def test_pool_avg():
pool = Pool("avg", 2, stride=2)
x = torch.randn(1, IN_CHANNELS, 64, 64)
out = pool(x)
assert out.shape == (1, IN_CHANNELS, 32, 32)
def test_adown():
adown = ADown(IN_CHANNELS, OUT_CHANNELS)
x = torch.randn(1, IN_CHANNELS, 64, 64)
out = adown(x)
assert out.shape == (1, OUT_CHANNELS, 32, 32)
def test_cblinear():
cblinear = CBLinear(IN_CHANNELS, [5, 5])
x = torch.randn(1, IN_CHANNELS, 64, 64)
outs = cblinear(x)
assert len(outs) == 2
assert outs[0].shape == (1, 5, 64, 64)
assert outs[1].shape == (1, 5, 64, 64)
def test_sppelan():
sppelan = SPPELAN(IN_CHANNELS, OUT_CHANNELS, NECK_CHANNELS)
x = torch.randn(1, IN_CHANNELS, 64, 64)
out = sppelan(x)
assert out.shape == (1, OUT_CHANNELS, 64, 64)
|