Upload subsampling.py with huggingface_hub

58f7f49 verified 5 months ago

5.74 kB

import math
import mlx.core as mx
import mlx.nn as nn

def calc_length(lengths, all_paddings, kernel_size, stride, repeat_num=1):
    """Calculates the output length of a Tensor passed through a convolution layer"""
    add_pad = all_paddings - kernel_size
    for i in range(repeat_num):
        lengths = mx.floor((lengths + add_pad) / stride) + 1
    return lengths.astype(mx.int32)

class ConvSubsampling(nn.Module):
    """
    MLX Translation of NVIDIA NeMo ConvSubsampling
    Used for the front-end of FastConformer.
    """
    def __init__(
        self,
        subsampling: str,
        subsampling_factor: int,
        feat_in: int,
        feat_out: int,
        conv_channels: int,
    ):
        super().__init__()
        self._subsampling = subsampling
        self._conv_channels = conv_channels
        self._feat_in = feat_in
        self._feat_out = feat_out
        
        self._sampling_num = int(math.log(subsampling_factor, 2))
        self.subsampling_factor = subsampling_factor
        
        in_channels = 1
        self.conv = []
        
        if subsampling == 'striding':
            self._stride = 2
            self._kernel_size = 3
            self._left_padding = (self._kernel_size - 1) // 2
            self._right_padding = (self._kernel_size - 1) // 2
            
            for i in range(self._sampling_num):
                self.conv.append(
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=conv_channels,
                        kernel_size=self._kernel_size,
                        stride=self._stride,
                        padding=self._left_padding,
                    )
                )
                in_channels = conv_channels
                
        elif subsampling == 'dw_striding':
            self._stride = 2
            self._kernel_size = 3
            self._left_padding = (self._kernel_size - 1) // 2
            self._right_padding = (self._kernel_size - 1) // 2
            
            # Layer 0: Full Conv
            self.conv.append(
                nn.Conv2d(
                    in_channels=in_channels,
                    out_channels=conv_channels,
                    kernel_size=self._kernel_size,
                    stride=self._stride,
                    padding=self._left_padding,
                )
            )
            # Layer 1: ReLU (not a module with weights, but we need to keep indexing aligned, so we use None or handle it differently)
            self.conv.append(None) # Index 1
            
            in_channels = conv_channels
            
            for i in range(self._sampling_num - 1):
                # Depthwise Conv (Layer 2, 5...)
                self.conv.append(
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=in_channels,
                        kernel_size=self._kernel_size,
                        stride=self._stride,
                        padding=self._left_padding,
                        groups=in_channels,
                    )
                )
                
                # Pointwise Conv (Layer 3, 6...)
                self.conv.append(
                    nn.Conv2d(
                        in_channels=in_channels,
                        out_channels=conv_channels,
                        kernel_size=1,
                        stride=1,
                        padding=0,
                        groups=1,
                    )
                )
                
                # ReLU placeholder (Layer 4, 7...)
                self.conv.append(None)
                in_channels = conv_channels
                
        else:
            raise ValueError(f"Subsampling {subsampling} not implemented in MLX yet")

        # In PyTorch: b, c, t, f -> b, t, c*f
        # In MLX: inputs to Conv2d are expected to be (N, H, W, C)
        # So we map: N=batch, H=time, W=freq, C=channels
        # After convs: flatten the last two dimensions to project to feat_out
        
        # Calculate output frequency size mathematically or dynamically
        dummy_f = mx.array([feat_in], dtype=mx.float32)
        out_f = calc_length(
            dummy_f,
            all_paddings=self._left_padding + self._right_padding,
            kernel_size=self._kernel_size,
            stride=self._stride,
            repeat_num=self._sampling_num
        )
        out_f_val = out_f.item()
        
        self.out = nn.Linear(conv_channels * int(out_f_val), feat_out)

    def __call__(self, x, lengths):
        """
        x: shape (batch, time, freq)
        lengths: shape (batch,)
        """
        # Convert to MLX Conv2d shape: (N, H, W, C) -> (batch, time, freq, 1)
        x = mx.expand_dims(x, axis=-1)
        
        for conv in self.conv:
            if conv is None:
                x = mx.maximum(x, 0.0) # ReLU
            else:
                x = conv(x)
            
        b, t, f, c = x.shape
        
        # In PyTorch, shape is (b, c, t, f) which is transposed to (b, t, c, f) and then flattened to (b, t, c*f).
        # In MLX, shape is (b, t, f, c). We MUST transpose to (b, t, c, f) before flattening to match the trained Linear weights.
        x = mx.transpose(x, (0, 1, 3, 2)) # shape becomes (b, t, c, f)
        
        # Flatten c and f
        x = mx.reshape(x, (b, t, c * f))
        
        x = self.out(x)
        
        out_lengths = calc_length(
            lengths,
            all_paddings=self._left_padding + self._right_padding,
            kernel_size=self._kernel_size,
            stride=self._stride,
            repeat_num=self._sampling_num
        )
        
        return x, out_lengths

	import math
	import mlx.core as mx
	import mlx.nn as nn

	def calc_length(lengths, all_paddings, kernel_size, stride, repeat_num=1):
	"""Calculates the output length of a Tensor passed through a convolution layer"""
	add_pad = all_paddings - kernel_size
	for i in range(repeat_num):
	lengths = mx.floor((lengths + add_pad) / stride) + 1
	return lengths.astype(mx.int32)

	class ConvSubsampling(nn.Module):
	"""
	MLX Translation of NVIDIA NeMo ConvSubsampling
	Used for the front-end of FastConformer.
	"""
	def __init__(
	self,
	subsampling: str,
	subsampling_factor: int,
	feat_in: int,
	feat_out: int,
	conv_channels: int,
	):
	super().__init__()
	self._subsampling = subsampling
	self._conv_channels = conv_channels
	self._feat_in = feat_in
	self._feat_out = feat_out

	self._sampling_num = int(math.log(subsampling_factor, 2))
	self.subsampling_factor = subsampling_factor

	in_channels = 1
	self.conv = []

	if subsampling == 'striding':
	self._stride = 2
	self._kernel_size = 3
	self._left_padding = (self._kernel_size - 1) // 2
	self._right_padding = (self._kernel_size - 1) // 2

	for i in range(self._sampling_num):
	self.conv.append(
	nn.Conv2d(
	in_channels=in_channels,
	out_channels=conv_channels,
	kernel_size=self._kernel_size,
	stride=self._stride,
	padding=self._left_padding,
	)
	)
	in_channels = conv_channels

	elif subsampling == 'dw_striding':
	self._stride = 2
	self._kernel_size = 3
	self._left_padding = (self._kernel_size - 1) // 2
	self._right_padding = (self._kernel_size - 1) // 2

	# Layer 0: Full Conv
	self.conv.append(
	nn.Conv2d(
	in_channels=in_channels,
	out_channels=conv_channels,
	kernel_size=self._kernel_size,
	stride=self._stride,
	padding=self._left_padding,
	)
	)
	# Layer 1: ReLU (not a module with weights, but we need to keep indexing aligned, so we use None or handle it differently)
	self.conv.append(None) # Index 1

	in_channels = conv_channels

	for i in range(self._sampling_num - 1):
	# Depthwise Conv (Layer 2, 5...)
	self.conv.append(
	nn.Conv2d(
	in_channels=in_channels,
	out_channels=in_channels,
	kernel_size=self._kernel_size,
	stride=self._stride,
	padding=self._left_padding,
	groups=in_channels,
	)
	)

	# Pointwise Conv (Layer 3, 6...)
	self.conv.append(
	nn.Conv2d(
	in_channels=in_channels,
	out_channels=conv_channels,
	kernel_size=1,
	stride=1,
	padding=0,
	groups=1,
	)
	)

	# ReLU placeholder (Layer 4, 7...)
	self.conv.append(None)
	in_channels = conv_channels

	else:
	raise ValueError(f"Subsampling {subsampling} not implemented in MLX yet")

	# In PyTorch: b, c, t, f -> b, t, c*f
	# In MLX: inputs to Conv2d are expected to be (N, H, W, C)
	# So we map: N=batch, H=time, W=freq, C=channels
	# After convs: flatten the last two dimensions to project to feat_out

	# Calculate output frequency size mathematically or dynamically
	dummy_f = mx.array([feat_in], dtype=mx.float32)
	out_f = calc_length(
	dummy_f,
	all_paddings=self._left_padding + self._right_padding,
	kernel_size=self._kernel_size,
	stride=self._stride,
	repeat_num=self._sampling_num
	)
	out_f_val = out_f.item()

	self.out = nn.Linear(conv_channels * int(out_f_val), feat_out)

	def __call__(self, x, lengths):
	"""
	x: shape (batch, time, freq)
	lengths: shape (batch,)
	"""
	# Convert to MLX Conv2d shape: (N, H, W, C) -> (batch, time, freq, 1)
	x = mx.expand_dims(x, axis=-1)

	for conv in self.conv:
	if conv is None:
	x = mx.maximum(x, 0.0) # ReLU
	else:
	x = conv(x)

	b, t, f, c = x.shape

	# In PyTorch, shape is (b, c, t, f) which is transposed to (b, t, c, f) and then flattened to (b, t, c*f).
	# In MLX, shape is (b, t, f, c). We MUST transpose to (b, t, c, f) before flattening to match the trained Linear weights.
	x = mx.transpose(x, (0, 1, 3, 2)) # shape becomes (b, t, c, f)

	# Flatten c and f
	x = mx.reshape(x, (b, t, c * f))

	x = self.out(x)

	out_lengths = calc_length(
	lengths,
	all_paddings=self._left_padding + self._right_padding,
	kernel_size=self._kernel_size,
	stride=self._stride,
	repeat_num=self._sampling_num
	)

	return x, out_lengths