|
|
|
|
|
|
|
|
import copy |
|
|
import unittest |
|
|
|
|
|
import torch |
|
|
import torch.nn as nn |
|
|
from pytorchvideo.layers.squeeze_excitation import ( |
|
|
create_audio_2d_squeeze_excitation_block, |
|
|
) |
|
|
|
|
|
|
|
|
class Test2DSqueezeExcitationBlock(unittest.TestCase): |
|
|
def setUp(self): |
|
|
|
|
|
self.layer_args = { |
|
|
"dim_in": 32, |
|
|
"dim_out": 32, |
|
|
"use_se": True, |
|
|
"se_reduction_ratio": 16, |
|
|
"branch_fusion": lambda x, y: x + y, |
|
|
"conv_a_kernel_size": 3, |
|
|
"conv_a_stride": 1, |
|
|
"conv_a_padding": 1, |
|
|
"conv_b_kernel_size": 3, |
|
|
"conv_b_stride": 1, |
|
|
"conv_b_padding": 1, |
|
|
"norm": nn.BatchNorm2d, |
|
|
"norm_eps": 1e-5, |
|
|
"norm_momentum": 0.1, |
|
|
"activation": nn.ReLU, |
|
|
} |
|
|
|
|
|
self.batchsize = 1 |
|
|
self.forward_pass_configs = [ |
|
|
{ |
|
|
"input": torch.rand(self.batchsize, self.layer_args["dim_in"], 100, 40), |
|
|
"output_shape": torch.Size( |
|
|
[self.batchsize, self.layer_args["dim_out"], 100, 40] |
|
|
), |
|
|
}, |
|
|
] |
|
|
|
|
|
def test_forward_pass(self): |
|
|
for split_config in self.forward_pass_configs: |
|
|
layer_args = copy.deepcopy(self.layer_args) |
|
|
model = create_audio_2d_squeeze_excitation_block(**layer_args) |
|
|
|
|
|
out = model(split_config["input"]) |
|
|
self.assertTrue(isinstance(out, torch.Tensor)) |
|
|
self.assertEqual(out.size(), split_config["output_shape"]) |
|
|
|
