MusicGen / tests /modules /test_conv.py
adefossez's picture
Initial commit
5238467
raw
history blame
7.38 kB
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
from itertools import product
import math
import random
import pytest
import torch
from torch import nn
from audiocraft.modules import (
NormConv1d,
NormConvTranspose1d,
StreamableConv1d,
StreamableConvTranspose1d,
pad1d,
unpad1d,
)
def test_get_extra_padding_for_conv1d():
# TODO: Implement me!
pass
def test_pad1d_zeros():
x = torch.randn(1, 1, 20)
xp1 = pad1d(x, (0, 5), mode='constant', value=0.)
assert xp1.shape[-1] == 25
xp2 = pad1d(x, (5, 5), mode='constant', value=0.)
assert xp2.shape[-1] == 30
xp3 = pad1d(x, (0, 0), mode='constant', value=0.)
assert xp3.shape[-1] == 20
xp4 = pad1d(x, (10, 30), mode='constant', value=0.)
assert xp4.shape[-1] == 60
with pytest.raises(AssertionError):
pad1d(x, (-1, 0), mode='constant', value=0.)
with pytest.raises(AssertionError):
pad1d(x, (0, -1), mode='constant', value=0.)
with pytest.raises(AssertionError):
pad1d(x, (-1, -1), mode='constant', value=0.)
def test_pad1d_reflect():
x = torch.randn(1, 1, 20)
xp1 = pad1d(x, (0, 5), mode='reflect', value=0.)
assert xp1.shape[-1] == 25
xp2 = pad1d(x, (5, 5), mode='reflect', value=0.)
assert xp2.shape[-1] == 30
xp3 = pad1d(x, (0, 0), mode='reflect', value=0.)
assert xp3.shape[-1] == 20
xp4 = pad1d(x, (10, 30), mode='reflect', value=0.)
assert xp4.shape[-1] == 60
with pytest.raises(AssertionError):
pad1d(x, (-1, 0), mode='reflect', value=0.)
with pytest.raises(AssertionError):
pad1d(x, (0, -1), mode='reflect', value=0.)
with pytest.raises(AssertionError):
pad1d(x, (-1, -1), mode='reflect', value=0.)
def test_unpad1d():
x = torch.randn(1, 1, 20)
u1 = unpad1d(x, (5, 5))
assert u1.shape[-1] == 10
u2 = unpad1d(x, (0, 5))
assert u2.shape[-1] == 15
u3 = unpad1d(x, (5, 0))
assert u3.shape[-1] == 15
u4 = unpad1d(x, (0, 0))
assert u4.shape[-1] == x.shape[-1]
with pytest.raises(AssertionError):
unpad1d(x, (-1, 0))
with pytest.raises(AssertionError):
unpad1d(x, (0, -1))
with pytest.raises(AssertionError):
unpad1d(x, (-1, -1))
class TestNormConv1d:
def test_norm_conv1d_modules(self):
N, C, T = 2, 2, random.randrange(1, 100_000)
t0 = torch.randn(N, C, T)
C_out, kernel_size, stride = 1, 4, 1
expected_out_length = int((T - kernel_size) / stride + 1)
wn_conv = NormConv1d(C, 1, kernel_size=4, norm='weight_norm')
gn_conv = NormConv1d(C, 1, kernel_size=4, norm='time_group_norm')
nn_conv = NormConv1d(C, 1, kernel_size=4, norm='none')
assert isinstance(wn_conv.norm, nn.Identity)
assert isinstance(wn_conv.conv, nn.Conv1d)
assert isinstance(gn_conv.norm, nn.GroupNorm)
assert isinstance(gn_conv.conv, nn.Conv1d)
assert isinstance(nn_conv.norm, nn.Identity)
assert isinstance(nn_conv.conv, nn.Conv1d)
for conv_layer in [wn_conv, gn_conv, nn_conv]:
out = conv_layer(t0)
assert isinstance(out, torch.Tensor)
assert list(out.shape) == [N, C_out, expected_out_length]
class TestNormConvTranspose1d:
def test_normalizations(self):
N, C, T = 2, 2, random.randrange(1, 100_000)
t0 = torch.randn(N, C, T)
C_out, kernel_size, stride = 1, 4, 1
expected_out_length = (T - 1) * stride + (kernel_size - 1) + 1
wn_convtr = NormConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride, norm='weight_norm')
gn_convtr = NormConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride, norm='time_group_norm')
nn_convtr = NormConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride, norm='none')
assert isinstance(wn_convtr.norm, nn.Identity)
assert isinstance(wn_convtr.convtr, nn.ConvTranspose1d)
assert isinstance(gn_convtr.norm, nn.GroupNorm)
assert isinstance(gn_convtr.convtr, nn.ConvTranspose1d)
assert isinstance(nn_convtr.norm, nn.Identity)
assert isinstance(nn_convtr.convtr, nn.ConvTranspose1d)
for convtr_layer in [wn_convtr, gn_convtr, nn_convtr]:
out = convtr_layer(t0)
assert isinstance(out, torch.Tensor)
assert list(out.shape) == [N, C_out, expected_out_length]
class TestStreamableConv1d:
def get_streamable_conv1d_output_length(self, length, kernel_size, stride, dilation):
# StreamableConv1d internally pads to make sure that the last window is full
padding_total = (kernel_size - 1) * dilation - (stride - 1)
n_frames = (length - kernel_size + padding_total) / stride + 1
ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
return ideal_length // stride
def test_streamable_conv1d(self):
N, C, T = 2, 2, random.randrange(1, 100_000)
t0 = torch.randn(N, C, T)
C_out = 1
# conv params are [(kernel_size, stride, dilation)]
conv_params = [(4, 1, 1), (4, 2, 1), (3, 1, 3), (10, 5, 1), (3, 2, 3)]
for causal, (kernel_size, stride, dilation) in product([False, True], conv_params):
expected_out_length = self.get_streamable_conv1d_output_length(T, kernel_size, stride, dilation)
sconv = StreamableConv1d(C, C_out, kernel_size=kernel_size, stride=stride, dilation=dilation, causal=causal)
out = sconv(t0)
assert isinstance(out, torch.Tensor)
print(list(out.shape), [N, C_out, expected_out_length])
assert list(out.shape) == [N, C_out, expected_out_length]
class TestStreamableConvTranspose1d:
def get_streamable_convtr1d_output_length(self, length, kernel_size, stride):
padding_total = (kernel_size - stride)
return (length - 1) * stride - padding_total + (kernel_size - 1) + 1
def test_streamable_convtr1d(self):
N, C, T = 2, 2, random.randrange(1, 100_000)
t0 = torch.randn(N, C, T)
C_out = 1
with pytest.raises(AssertionError):
StreamableConvTranspose1d(C, C_out, kernel_size=4, causal=False, trim_right_ratio=0.5)
StreamableConvTranspose1d(C, C_out, kernel_size=4, causal=True, trim_right_ratio=-1.)
StreamableConvTranspose1d(C, C_out, kernel_size=4, causal=True, trim_right_ratio=2)
# causal params are [(causal, trim_right)]
causal_params = [(False, 1.0), (True, 1.0), (True, 0.5), (True, 0.0)]
# conv params are [(kernel_size, stride)]
conv_params = [(4, 1), (4, 2), (3, 1), (10, 5)]
for ((causal, trim_right_ratio), (kernel_size, stride)) in product(causal_params, conv_params):
expected_out_length = self.get_streamable_convtr1d_output_length(T, kernel_size, stride)
sconvtr = StreamableConvTranspose1d(C, C_out, kernel_size=kernel_size, stride=stride,
causal=causal, trim_right_ratio=trim_right_ratio)
out = sconvtr(t0)
assert isinstance(out, torch.Tensor)
assert list(out.shape) == [N, C_out, expected_out_length]