MusicGen

Runtime error

App Files Files Community

MusicGen / tests /modules /test_seanet.py

adefossez

Initial commit

5238467 about 1 year ago

raw history blame contribute delete

No virus

4.87 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 pytest
	import torch

	from audiocraft.modules.seanet import SEANetEncoder, SEANetDecoder, SEANetResnetBlock
	from audiocraft.modules import StreamableConv1d, StreamableConvTranspose1d


	class TestSEANetModel:

	def test_base(self):
	encoder = SEANetEncoder()
	decoder = SEANetDecoder()

	x = torch.randn(1, 1, 24000)
	z = encoder(x)
	assert list(z.shape) == [1, 128, 75], z.shape
	y = decoder(z)
	assert y.shape == x.shape, (x.shape, y.shape)

	def test_causal(self):
	encoder = SEANetEncoder(causal=True)
	decoder = SEANetDecoder(causal=True)
	x = torch.randn(1, 1, 24000)

	z = encoder(x)
	assert list(z.shape) == [1, 128, 75], z.shape
	y = decoder(z)
	assert y.shape == x.shape, (x.shape, y.shape)

	def test_conv_skip_connection(self):
	encoder = SEANetEncoder(true_skip=False)
	decoder = SEANetDecoder(true_skip=False)

	x = torch.randn(1, 1, 24000)
	z = encoder(x)
	assert list(z.shape) == [1, 128, 75], z.shape
	y = decoder(z)
	assert y.shape == x.shape, (x.shape, y.shape)

	def test_seanet_encoder_decoder_final_act(self):
	encoder = SEANetEncoder(true_skip=False)
	decoder = SEANetDecoder(true_skip=False, final_activation='Tanh')

	x = torch.randn(1, 1, 24000)
	z = encoder(x)
	assert list(z.shape) == [1, 128, 75], z.shape
	y = decoder(z)
	assert y.shape == x.shape, (x.shape, y.shape)

	def _check_encoder_blocks_norm(self, encoder: SEANetEncoder, n_disable_blocks: int, norm: str):
	n_blocks = 0
	for layer in encoder.model:
	if isinstance(layer, StreamableConv1d):
	n_blocks += 1
	assert layer.conv.norm_type == 'none' if n_blocks <= n_disable_blocks else norm
	elif isinstance(layer, SEANetResnetBlock):
	for resnet_layer in layer.block:
	if isinstance(resnet_layer, StreamableConv1d):
	# here we add + 1 to n_blocks as we increment n_blocks just after the block
	assert resnet_layer.conv.norm_type == 'none' if (n_blocks + 1) <= n_disable_blocks else norm

	def test_encoder_disable_norm(self):
	n_residuals = [0, 1, 3]
	disable_blocks = [0, 1, 2, 3, 4, 5, 6]
	norms = ['weight_norm', 'none']
	for n_res, disable_blocks, norm in product(n_residuals, disable_blocks, norms):
	encoder = SEANetEncoder(n_residual_layers=n_res, norm=norm,
	disable_norm_outer_blocks=disable_blocks)
	self._check_encoder_blocks_norm(encoder, disable_blocks, norm)

	def _check_decoder_blocks_norm(self, decoder: SEANetDecoder, n_disable_blocks: int, norm: str):
	n_blocks = 0
	for layer in decoder.model:
	if isinstance(layer, StreamableConv1d):
	n_blocks += 1
	assert layer.conv.norm_type == 'none' if (decoder.n_blocks - n_blocks) < n_disable_blocks else norm
	elif isinstance(layer, StreamableConvTranspose1d):
	n_blocks += 1
	assert layer.convtr.norm_type == 'none' if (decoder.n_blocks - n_blocks) < n_disable_blocks else norm
	elif isinstance(layer, SEANetResnetBlock):
	for resnet_layer in layer.block:
	if isinstance(resnet_layer, StreamableConv1d):
	assert resnet_layer.conv.norm_type == 'none' \
	if (decoder.n_blocks - n_blocks) < n_disable_blocks else norm

	def test_decoder_disable_norm(self):
	n_residuals = [0, 1, 3]
	disable_blocks = [0, 1, 2, 3, 4, 5, 6]
	norms = ['weight_norm', 'none']
	for n_res, disable_blocks, norm in product(n_residuals, disable_blocks, norms):
	decoder = SEANetDecoder(n_residual_layers=n_res, norm=norm,
	disable_norm_outer_blocks=disable_blocks)
	self._check_decoder_blocks_norm(decoder, disable_blocks, norm)

	def test_disable_norm_raises_exception(self):
	# Invalid disable_norm_outer_blocks values raise exceptions
	with pytest.raises(AssertionError):
	SEANetEncoder(disable_norm_outer_blocks=-1)

	with pytest.raises(AssertionError):
	SEANetEncoder(ratios=[1, 1, 2, 2], disable_norm_outer_blocks=7)

	with pytest.raises(AssertionError):
	SEANetDecoder(disable_norm_outer_blocks=-1)

	with pytest.raises(AssertionError):
	SEANetDecoder(ratios=[1, 1, 2, 2], disable_norm_outer_blocks=7)