File size: 4,056 Bytes
3133fdb |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import itertools
import os
import unittest
import numpy as np
import torch
from pytorchvideo.models.csn import create_csn
from pytorchvideo.models.resnet import create_bottleneck_block
from torch import nn
class TestCSN(unittest.TestCase):
def setUp(self):
super().setUp()
torch.set_rng_state(torch.manual_seed(42).get_state())
def test_create_csn(self):
"""
Test simple CSN with different inputs.
"""
for input_channel, input_clip_length, input_crop_size in itertools.product(
(3, 2), (4, 8), (56, 64)
):
stage_spatial_stride = (1, 2, 2, 2)
stage_temporal_stride = (1, 2, 2, 1)
total_spatial_stride = 2 * np.prod(stage_spatial_stride)
total_temporal_stride = np.prod(stage_temporal_stride)
head_pool_kernel_size = (
input_clip_length // total_temporal_stride,
input_crop_size // total_spatial_stride,
input_crop_size // total_spatial_stride,
)
model = create_csn(
input_channel=input_channel,
model_depth=50,
model_num_class=400,
dropout_rate=0,
norm=nn.BatchNorm3d,
activation=nn.ReLU,
stem_dim_out=8,
stem_conv_kernel_size=(3, 7, 7),
stem_conv_stride=(1, 2, 2),
stage_conv_a_kernel_size=(1, 1, 1),
stage_conv_b_kernel_size=(3, 3, 3),
stage_conv_b_width_per_group=1,
stage_spatial_stride=(1, 2, 2, 2),
stage_temporal_stride=(1, 2, 2, 1),
bottleneck=create_bottleneck_block,
head_pool=nn.AvgPool3d,
head_pool_kernel_size=head_pool_kernel_size,
head_output_size=(1, 1, 1),
head_activation=nn.Softmax,
)
# Test forwarding.
for tensor in TestCSN._get_inputs(
input_channel, input_clip_length, input_crop_size
):
if tensor.shape[1] != input_channel:
with self.assertRaises(RuntimeError):
out = model(tensor)
continue
out = model(tensor)
output_shape = out.shape
output_shape_gt = (tensor.shape[0], 400)
self.assertEqual(
output_shape,
output_shape_gt,
"Output shape {} is different from expected shape {}".format(
output_shape, output_shape_gt
),
)
@staticmethod
def _get_inputs(
channel: int = 3, clip_length: int = 4, crop_size: int = 112
) -> torch.tensor:
"""
Provide different tensors as test cases.
Yield:
(torch.tensor): tensor as test case input.
"""
# Prepare random inputs as test cases.
shapes = (
(1, channel, clip_length, crop_size, crop_size),
(2, channel, clip_length, crop_size, crop_size),
)
for shape in shapes:
yield torch.rand(shape)
def test_load_hubconf(self):
path = os.path.join(
os.path.dirname(os.path.realpath(__file__)),
"..",
)
input_channel = 3
input_clip_length = 4
input_crop_size = 56
model = torch.hub.load(
repo_or_dir=path, source="local", model="csn_r101", pretrained=False
)
self.assertIsNotNone(model)
# Test forwarding.
for tensor in TestCSN._get_inputs(
input_channel, input_clip_length, input_crop_size
):
with torch.no_grad():
if tensor.shape[1] != input_channel:
with self.assertRaises(RuntimeError):
model(tensor)
continue
|