Diffusion models with Adaptive Guidance for 16-electrode Electrical Impedance Tomography (DAG-EIT)

Implementation of DAG-EIT in PyTorch.

Installation

Clone the repository and move to the project directory:

git clone <current url>
cd DAG-EIT

Create the Conda environment using the provided environment.yml file:

conda env create -f environment.yml

Activate the environment:

conda activate DAG-EIT

To verify that PyTorch and CUDA are available, run:

python -c "import torch; print('PyTorch version:', torch.__version__); print('CUDA available:', torch.cuda.is_available())"

Data and Pretrained Model

The datasets and pretrained model can be downloaded from the following links:

Resource	Download Link
Training Data	Google Drive
Validation Data	Google Drive
Test Data	Google Drive
Pretrained Model	Google Drive
All Datasets and Pretrained Model	Google Drive

Place the downloaded training and validation files in the following directories:

DAG-EIT/
├── data/
│   ├── training/
│   │   ├── eit_train_0326_original.pt
│   │   ├── eit_train_0326_obesity1.pt
│   │   ├── eit_train_0326_obesity2.pt
│   │   └── eit_train_0326_ellipse.pt
│   ├── validation/
│   │   └── eit_valid_0326.pt
│   └── testing/
│       └── eit_test_0326.pt
└── pretrained/best_model.pkl

Before sampling, update the test-data path and pretrained-model path in configs/eit.yaml.

Training

Set YOUR_DEVICE_NUMBER to the GPU device number to use. For example, use DEVICE=0 for the first GPU.

DEVICE=YOUR_DEVICE_NUMBER

nohup bash -c "CUDA_VISIBLE_DEVICES=${DEVICE} python train.py \
  --outdir=./pretrained \
  --data=./data/training/eit_train_0326_original.pt,./data/training/eit_train_0326_obesity1.pt,./data/training/eit_train_0326_obesity2.pt,./data/training/eit_train_0326_ellipse.pt \
  --train-mix-ratio=2,1,1,1 \
  --valid-data=./data/validation/eit_valid_0326.pt \
  --cond=0 --arch=ddpmpp \
  --batch=100 --batch-gpu=100 \
  --tick=10 --snap=100 --dump=500 --duration=30 --ema=0.05 \
  --lr=1e-4" \
  > eit_train.txt 2>&1 &

The training process runs in the background with nohup, and the output is written to eit_train.txt.

To monitor the training log, run:

tail -f eit_train.txt

Sampling

Configure the test dataset, pretrained model, and GPU device in configs/eit.yaml, and then run:

python generate_eit.py --config configs/eit.yaml

Citation

:::