FunASR/funasr/models/preencoder/linear.py
jmwang66 98abc0e5ac
update setup (#686)
* update

* update setup

* update setup

* update setup

* update setup

* update setup

* update setup

* update

* update

* update setup
2023-06-29 16:30:39 +08:00

37 lines
961 B
Python

#!/usr/bin/env python3
# 2021, Carnegie Mellon University; Xuankai Chang
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
"""Linear Projection."""
from funasr.models.preencoder.abs_preencoder import AbsPreEncoder
from typing import Tuple
import torch
class LinearProjection(AbsPreEncoder):
"""Linear Projection Preencoder."""
def __init__(
self,
input_size: int,
output_size: int,
):
"""Initialize the module."""
super().__init__()
self.output_dim = output_size
self.linear_out = torch.nn.Linear(input_size, output_size)
def forward(
self, input: torch.Tensor, input_lengths: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
"""Forward."""
output = self.linear_out(input)
return output, input_lengths # no state in this layer
def output_size(self) -> int:
"""Get the output size."""
return self.output_dim