TRANSFORMER COMPRESSED SENSING VIA GLOBAL IMAGE TOKENS

=========================================================================

This repository contains the supplementary material and PyTorch implementation of Deep cascade of Transformer Neural Networks (DcTNN).

Environment Dependencies

Code is known to work with:

• Python 3.8.10
• PyTorch 1.9.0
• Einops 0.3.0
• Pillow 8.0.0

---

Usage

build_model.py provides an example of how to build and test DcTNN, cascading Axial (Ax), Kaleidoscope (KD) and Patch (P) transformer encoder layers to build our Ensemble network.

Class cascadeNet takes the following arguments:

+---------------+-----------+------------------------------------------------------------------+ | Args | Type | Description | +===============+===========+==================================================================+ | N | int | Image Size | +---------------+-----------+------------------------------------------------------------------+ | encList | array | Is an array that contains denoising transformer encoder modules | +---------------+-----------+------------------------------------------------------------------+ | encArgs | array | Contains dictionaries with args for encoders in encList | +---------------+-----------+------------------------------------------------------------------+ | dcFunc | function | Contains the data consistency function to be used in recon | +---------------+-----------+------------------------------------------------------------------+ | lamb | bool | Whether or not to use a leanred data consistency parameter | +---------------+-----------+------------------------------------------------------------------+

Classes axVIT and patchVIT construct image denoisers based on axialEncoder and imageEncoder respectively; patchVIT can be configured as a Patch or Kaleidoscope Encoder.

Therefore to build the denoiser featured in our paper:

.. code-block:: python

Import necessary models

from DcTNN.tnn import * from dc.dc import *

Network parameters

Image size

N = 320

Size of patches/kd tokens

patchSize = 16

Number of heads for patch/kd encoders

nhead_patch = 8

Number of heads for axial encoder

nhead_axial = 8

Number of cascaded denoising blocks within each TNN

layerNo = 1

None d_model defaults to input dimension

d_model_axial = None d_model_patch = None

Number of encoder layers per-transformer

num_encoder_layers = 2

Number of channels of the image

numCh = 1

None dim_feedforward defaults to d_model^(3/2)

dim_feedforward = None

Define the dictionaries of parameter values

patchArgs = {"patch_size": patchSize, "kaleidoscope": False, "layerNo": layerNo, "numCh": numCh, "nhead": nhead_patch, "num_encoder_layers": num_encoder_layers, "dim_feedforward": dim_feedforward, "d_model": d_model_patch} kdArgs = {"patch_size": patchSize, "kaleidoscope": True, "layerNo": layerNo, "numCh": numCh, "nhead": nhead_patch, "num_encoder_layers": num_encoder_layers, "dim_feedforward": dim_feedforward, "d_model": d_model_patch} axArgs = {"layerNo": layerNo, "numCh": numCh, "d_model": d_model_axial, "nhead": nhead_axial, "num_encoder_layers": num_encoder_layers, "dim_feedforward": dim_feedforward}

Build the list of encoders being used

encList = [axVIT, patchVIT, patchVIT]

Arguments to feed into encoders

encArgs = [axArgs, kdArgs, patchArgs]

Data consistency function

dcFunc = FFT_DC

Use learned weighting parameter

lamb = True

Define the model

dcenc = cascadeNet(N, encList, encArgs, dcFunc, lamb)

Citation and Acknowledgement

Paper is available on IEEE Xplore https://doi.org/10.1109/ICIP46576.2022.9897630::

M. B. Lorenzana, C. Engstrom and S. S. Chandra, "Transformer Compressed Sensing Via Global Image Tokens," 2022 IEEE International Conference on Image Processing (ICIP), 2022, pp. 3011-3015, doi: 10.1109/ICIP46576.2022.9897630.

---

Paper is also available on arXiv https://arxiv.org/pdf/2203.12861.pdf::

M. Bran Lorenzana, C. Engstrom, and S. S. Chandra, ‘Transformer Compressed Sensing via Global Image Tokens’. arXiv, Mar. 27, 2022. Available: http://arxiv.org/abs/2203.12861

---

Kaleidoscope transform was introduced by White et. al https://doi.org/10.1109/LSP.2021.3116510::

J. M. White, S. Crozier and S. S. Chandra, "Bespoke Fractal Sampling Patterns for Discrete Fourier Space via the Kaleidoscope Transform," in IEEE Signal Processing Letters, vol. 28, pp. 2053-2057, 2021, doi: 10.1109/LSP.2021.3116510.

---