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:
## 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.