A PyTorch implementation of FNet from the paper FNet: Mixing Tokens with Fourier Transforms by James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon (arXiv).

Additional to the architecture implementation, this repository offers a script for converting a checkpoint from the official FNet implementation (written in Jax) to a PyTorch checkpoint (statedict and or model export).

We offer the following converted checkpoints and pre-trained models

The PyTorch checkpoints marked with converted are converted Jax checkpoints using the technique described below.

You can install this repository as a package running

pip install fnet-pytorch

Now, you can load a pre-trained model in PyTorch as follows. You'll need the config.json and the .statedict.pt file.

import torch
import json
from fnet import FNet, FNetForPretraining

with open('path/to/config.json', 'r') as f:
    config = json.load(f)

# if you just want the encoder
fnet = FNet(config)
fnet.load_state_dict(torch.load('path/to/fnet.statedict.pt'))

# if you want FNet with pre-training head
fnet = FNetForPretraining(config)
fnet.load_state_dict(torch.load('path/to/fnet_pretraining.statedict.pt'))

You can also get the config from only the state dict:

from fnet import get_config_from_statedict

state_dict = torch.load('path/to/fnet.statedict.pt')
config = get_config_from_statedict(state_dict)
fnet = FNet(config)
fnet.load_state_dict(state_dict)

But not all config values can be inferred from the state dict alone, like dropout rate, fourier layer type and padding token index. get_config_from_statedict uses reasonable defaults for them. Look into the implementation to see which parameters are not inferred and how it might affect your use case.

Download a pre-trained Jax checkpoint of FNet from their official GitHub page or use any checkpoint that you trained using the official implementation.
You also need the SentencePiece vocab model. For the official checkpoints, use the model given here. For custom checkpoints use your respective vocab model.

Install dependencies (ideally in a virtualenv)

pip install -r requirements.txt

Convert a checkpoint to PyTorch

python convert_jax_checkpoint.py \
    --checkpoint <path/to/checkout> \
    --vocab <path/to/vocab> \
    --outdir <outdir>

Output files: config.json, fnet.statedict.pt, fnet_pretraining.statedict.pt

The checkpoints from the official Jax implementation are of complete pre-training models, meaning they contain encoder and pre-training head weights. The conversion will convert the Jax checkpoint to a PyTorch statedict of this project's FNet module (fnet.statedict.pt) and FNetForPreTraining module (fnet_pretraining.statedict.pt). You can use the model type for your needs whether you want to run further pre-trainings or not.

Although all model parameters will be correctly transferred to the PyTorch model, there will be slight differences between Jax and PyTorch in the inference result because their LayerNorm and GELU implementations slightly differ.

For a given inference input, all hidden states and logits of the official and converted model are equal at least up the first digit after the comma. This is programmatically verified using the script described below.

You can use the verify_conversion.py script to compare the inference outputs of a Jax checkpoint vs. the converted PyTorch checkpoint. But since this requires properly running the Jax FNet it requires a bit of setup and some modifications to the official implementation.

1. Clone the official implementation

svn export https://github.com/google-research/google-research/trunk/f_net
# or
git clone git@github.com:google-research/google-research.git
cd google-research/f_net

2. Edit the config in the official implementation to fit the checkpoint you want to run.

3. Add the following to the return value of _compute_pretraining_metrics in models.py:

return {
    ...
    "masked_lm_logits": masked_lm_logits,
    "next_sentence_logits": next_sentence_logits
}

2. Create a setup.py file in the parent directory of the f_net directory with the following content

from setuptools import setup

setup(
    name='fnet_jax',
    version='0.1.0',
    install_requires=[],
    packages=['f_net']
)

3. Install as a dependency in your fnet-pytorch project

pip install -e path/to/dir-of-"setup.py"

python verify_conversion.py \
    --jax path/to/jax_checkpoint \
    --torch path/to/fnet_for_pretraining.statedict.pt \
    --config path/to/config.json \
    --vocab path/to/vocab

This should initialize both models from the checkpoints and run inference on a sample text and compare the output logits.