search

rwth-i6 / pytorch-to-returnn

Make PyTorch code runnable within RETURNN
3 stars 6 forks source link

readme

Make PyTorch <https://pytorch.org/> code runnable within RETURNN <https://github.com/rwth-i6/returnn> (on TensorFlow). This provides some wrappers (and maybe some magic) to do that.

Installation

This package is on PyPI <https://pypi.org/project/pytorch-to-returnn/>__.

.. code-block::

pip install pytorch-to-returnn

torch drop-in replacement for RETURNN

The idea:

.. code-block:: python

import torch

class Model(torch.nn.Module):
 ...

Can be changed to:

.. code-block:: python

from pytorch_to_returnn import torch as torch_returnn

class Model(torch_returnn.nn.Module):
 ...

And this can be used directly in RETURNN.

This would convert the model to a RETURNN model. Example constructed RETURNN net dict <https://gist.github.com/albertz/01264cfbd2dfd73a19c1e2ac40bdb16b>, created from this PyTorch code <https://github.com/albertz/import-parallel-wavegan/blob/main/pytorch_to_returnn.py>.

Why

From PyTorch perspective:

From RETURNN/TF perspective:

How does this work

On a high level, RETURNN layers mostly corresponds to PyTorch modules. So all PyTorch modules are mapped directly or indirectly to RETURNN layers. The same is done for all functions in functional.

All RETURNN layers have further meta information about tensors, esp their axes/dimensions, and they might reorder axes when this is more efficient. We keep track of the axis mapping.

See the documentation of the pytorch_to_returnn.torch package <pytorch_to_returnn/torch> for details about how this works, and what can be done with it. Obviously, this is incomplete. For some status of what is not supported currently, see the supported document <Supported.md> and the unsupported document <Unsupported.md>__. Otherwise, when you hit some Module or functional function, or Tensor function which is not implemented, it just means that no-one has implemented it yet.

Direct use in RETURNN

A RETURNN config could be written in this way.

Use some PyTorch model as a component / subnetwork:

.. code-block:: python

from pytorch_to_returnn import torch as torch_returnn

class MyTorchModel(torch_returnn.nn.Module):
  ...

my_torch_model = MyTorchModel()

extern_data = {...}  # as usual

# RETURNN network dict
network = {
"prenet": my_torch_model.as_returnn_layer_dict(extern_data["data"]),

# Other RETURNN layers
...
}

Or directly using a PyTorch model as-is:

.. code-block:: python

from pytorch_to_returnn import torch as torch_returnn

class MyTorchModel(torch_returnn.nn.Module):
  ...

my_torch_model = MyTorchModel()

extern_data = {...}  # as usual

# RETURNN network dict
network = my_torch_model.as_returnn_net_dict(extern_data["data"])

Model converter

For the process of converting a model from PyTorch to RETURNN, including a PyTorch model checkpoint, we provide some utilities to automate this, and verify whether all outputs match. This is in pytorch_to_returnn.converter <pytorch_to_returnn/converter>__.

Example for Parallel WaveGAN <https://github.com/kan-bayashi/ParallelWaveGAN>__:

.. code-block:: python

def model_func(wrapped_import, inputs: torch.Tensor):
    if typing.TYPE_CHECKING or not wrapped_import:
        import torch
        from parallel_wavegan import models as pwg_models
        from parallel_wavegan import layers as pwg_layers

    else:
        torch = wrapped_import("torch")
        wrapped_import("parallel_wavegan")
        pwg_models = wrapped_import("parallel_wavegan.models")
        pwg_layers = wrapped_import("parallel_wavegan.layers")

    # Initialize PWG
    pwg_config = yaml.load(open(args.pwg_config), Loader=yaml.Loader)
    generator = pwg_models.MelGANGenerator(**pwg_config['generator_params'])
    generator.load_state_dict(
        torch.load(args.pwg_checkpoint, map_location="cpu")["model"]["generator"])
    generator.remove_weight_norm()
    pwg_model = generator.eval()
    pwg_pqmf = pwg_layers.PQMF(pwg_config["generator_params"]["out_channels"])

    return pwg_pqmf.synthesis(pwg_model(inputs))

feature_data = numpy.load(args.features)  # shape (Batch,Channel,Time) (1,80,80)

from pytorch_to_returnn.converter import verify_torch_and_convert_to_returnn
verify_torch_and_convert_to_returnn(model_func, inputs=feature_data)

The wrapped_import uses some import wrappers, which automatically converts the import torch statements.

This will automatically do the conversion, i.e. create a RETURNN model, including the RETURNN net dict <https://gist.github.com/albertz/01264cfbd2dfd73a19c1e2ac40bdb16b>__ and TF checkpoint file, and do verification on several steps of all the outputs (PyTorch module outputs vs RETURNN layer outputs).

Import wrapper

We also support to transform external PyTorch code on-the-fly (without the need to rewrite the code; it translates the code on AST level in the way above on-the-fly). I.e. it basically replaces import torch by from pytorch_to_returnn import torch -- that's all it does.

This is via our generic Python import wrapper pytorch_to_returnn.import_wrapper <pytorch_to_returnn/import_wrapper>__.

Example for Parallel WaveGAN <https://github.com/kan-bayashi/ParallelWaveGAN>__:

.. code-block:: python

import tensorflow as tf
from pytorch_to_returnn.import_wrapper import wrapped_import_torch_returnn
from pytorch_to_returnn.naming import Naming
from returnn.tf.util.data import Data

torch = wrapped_import_torch_returnn("torch")
wrapped_import_torch_returnn("parallel_wavegan")
pwg_models = wrapped_import_torch_returnn("parallel_wavegan.models")
pwg_layers = wrapped_import_torch_returnn("parallel_wavegan.layers")

naming = Naming.get_instance()  # default instance

inputs = torch.from_numpy(inputs)  # shape (Batch,Channel,Time), e.g. (1,80,80)
x = naming.register_input(
    inputs, Data("data", shape=(80, None), feature_dim_axis=1, time_dim_axis=2))
assert isinstance(x, Data)

# Initialize PWG
pwg_config = yaml.load(open(args.pwg_config), Loader=yaml.Loader)
generator = pwg_models.MelGANGenerator(**pwg_config['generator_params'])
generator.load_state_dict(
    torch.load(args.pwg_checkpoint, map_location="cpu")["model"]["generator"])
generator.remove_weight_norm()
pwg_model = generator.eval()
pwg_pqmf = pwg_layers.PQMF(pwg_config["generator_params"]["out_channels"])

outputs = pwg_pqmf.synthesis(pwg_model(inputs))

outputs = naming.register_output(outputs)
y = outputs.returnn_data
assert isinstance(y, Data)
assert isinstance(y.placeholder, tf.Tensor)

(RETURNN Data encapsulates a tensor and adds a lot of meta information about it and its axes, such as sequence lengths, beam, vocabulary of class indices, etc.)

Examples

See examples <examples>__.

Tests

See tests <tests>__. They are automatically run via GitHub Actions for CI.

.. image:: https://github.com/rwth-i6/pytorch-to-returnn/workflows/CI/badge.svg :target: https://github.com/rwth-i6/pytorch-to-returnn/actions

Related work