train.sh

#!/bin/bash

echo "Start to train the model...."


dataroot="/Data/dataset/HDR/Kalantari17"  # set the path of dataset
ckpt="./ckpt"
gpu="1"   # set the id of GPUs
name="ahdrnet_stage1"   # set the name of the folder where the model is stored
net="AHDRNet"  # can only choose 'AHDRNet', 'FSHDR', 'HDR-Transformer', or 'SCTNet'

batch="16"
niter="150"
lr_decay="50"
lr="0.0001"

build_dir=$ckpt"/"$name
if [ ! -d "$build_dir" ]; then
        mkdir $build_dir
fi
LOG=$build_dir/`date +%Y-%m-%d-%H-%M-%S`.txt

# Step1:  Training 1-Stage model
python train.py \
    --dataset_name sig17align   --model selfhdr1       --name $name           --network $net    \
    --dataroot $dataroot        --patch_size 128       --print_freq 300       --gpu_ids $gpu    \
    --checkpoints_dir $ckpt     --batch_size $batch    --niter $niter         --lr_decay_iters $lr_decay    \
    --lr $lr                    --lr_policy step       --save_imgs False       | tee $LOG    




# Step2:  Inference the 1-stage model on training dataset
iter=$niter  # set the epoch number of the 1-stage model
python test.py \
    --dataset_name sig17align   --model selfhdr2     --name $name        --network $net    --dataroot $dataroot   \
    --checkpoints_dir $ckpt     --load_iter $iter    --save_imgs True    --gpu_ids $gpu    --stage0_inference True  




# Step3: Training 2-Stage model
stru_lable=$ckpt"/"$name"/output_"$iter   # set the path of structure label generated by 1-stage model
name="ahdrnet_stage2" 

build_dir=$ckpt"/"$name
if [ ! -d "$build_dir" ]; then
        mkdir $build_dir
fi
LOG=$build_dir/`date +%Y-%m-%d-%H-%M-%S`.txt

python train.py \
    --dataset_name sig17align   --model selfhdr2       --name $name         --network $net      \
    --dataroot $dataroot        --patch_size 128       --print_freq 300     --gpu_ids $gpu         \
    --checkpoints_dir $ckpt     --batch_size $batch    --niter $niter       --lr_decay_iters $lr_decay    \
    --lr $lr                    --lr_policy step       --save_imgs False    --stru_lable_path $stru_lable   | tee $LOG