release v1

Author: Xiaoyan-xyy

README.md

@@ -1,14 +1,14 @@
 # TW-autoencode - Tied-weights AutoEncoder
 
-This repository contains the source code of unrolled LRP model and baselines from the paper, **Model Guidance via Explanations Turns Image
-Classifiers into Segmentation Models**.
-
-- TODO add link for the paper.
+This repository contains the source code of unrolled LRP model and baselines from the paper, [**Model Guidance via Explanations Turns Image
+Classifiers into Segmentation Models**](https://arxiv.org/pdf/2407.03009).
 
 
 ## Updates
 
-**Apr. 2024** -- Clean the code.
+**Apr. 2024** -- Start to clean the code.
+
+**July. 2024** -- Publish the first version of code.
 
 ## Datasets
 
@@ -31,58 +31,45 @@ data_parent_folder
 - Supervised training for UNets in 4 different cases (20, 100, 500, 1464 pixel-labeled data)
 
 ```
-qsub train.sh train.py --model std_unet --data_path  /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 15000 --num_labels 20 --uniform_masks 1.0 --save_folder std_unet_resnet50_lab20_lr1e-5_s40
+python train.py --model std_unet --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 15000 --num_labels 20 --uniform_masks 1.0 --save_folder std_unet_resnet50_lab20_lr1e-5_s42
 ```
 
 ```
-qsub train.sh train.py --model std_unet --data_path  /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 3000 --num_labels 100 --uniform_masks 1.0 --save_folder std_unet_resnet50_lab100_lr1e-5_s40
+python train.py --model std_unet --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 3000 --num_labels 100 --uniform_masks 1.0 --save_folder std_unet_resnet50_lab100_lr1e-5_s42
 ```
 
 ```
-qsub train.sh train.py --model std_unet --data_path  /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 600 --num_labels 500 --uniform_masks 0.5 --save_folder std_unet_resnet50_lab500_lr1e-5_s40
+python train.py --model std_unet --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 600 --num_labels 500 --uniform_masks 0.5 --save_folder std_unet_resnet50_lab500_lr1e-5_s42
 ```
 
 ```
-qsub train.sh train.py --model std_unet --data_path  /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 200 --save_folder std_unet_resnet50_lab1464_lr1e-5_s40
+python train.py --model std_unet --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 200 --save_folder std_unet_resnet50_lab1464_lr1e-5_s42
 ```
 - Semi-supervised training for multi-task UNets and Unrolled LRP models in 4 different cases (20, 100, 500, 1464 pixel-labeled data). For multi-task UNets, change  `--model unrolled_lrp` to `--model mt_unet`.
 
 
 ```
-qsub train.sh train.py --model unrolled_lrp --semisup_dataset --data_path /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 100 --iterative_gradients --add_classification --num_labels 20 --uniform_masks 1.0 --save_folder lrp0_resnet50_lab20_lr1e-5_s40 
+python train.py --model unrolled_lrp --semisup_dataset --add_classification --iterative_gradients --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 100 --num_labels 20 --uniform_masks 1.0 --save_folder lrp0_resnet50_lab20_lr1e-5_s42
 ```
 
 ```
-qsub train.sh train.py --model unrolled_lrp --semisup_dataset --data_path /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 100 --iterative_gradients --add_classification --num_labels 100 --uniform_masks 1.0 --save_folder lrp0_resnet50_lab100_lr1e-5_s40 
+python train.py --model unrolled_lrp --semisup_dataset --add_classification --iterative_gradients --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 100 --num_labels 100 --uniform_masks 1.0 --save_folder lrp0_resnet50_lab100_lr1e-5_s42 
 ```
 
 ```
-qsub train.sh train.py --model unrolled_lrp --semisup_dataset --data_path /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 100 --iterative_gradients --add_classification --num_labels 500 --uniform_masks 0.5 --save_folder lrp0_resnet50_lab500_lr1e-5_s40 
+python train.py --model unrolled_lrp --semisup_dataset --add_classification --iterative_gradients --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 100 --num_labels 500 --uniform_masks 0.5 --save_folder lrp0_resnet50_lab500_lr1e-5_s42
 ```
 
 ```
-qsub train.sh train.py --model unrolled_lrp --semisup_dataset --data_path /fast/AG_Kainmueller/xyu/ --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 40 --epochs 100 --iterative_gradients --add_classification --save_folder lrp0_resnet50_lab1464_lr1e-5_s40 
-```
-
-
-- TODO rewrite above comments which is used for running on the max cluster
-
-example
-```
-python train.py --model std_unet --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 15000 --num_labels 20 --uniform_masks 1.0 --save_folder std_unet_resnet50_lab20_lr1e-5_s42
-```
-
-
-```
-python train.py --model unrolled_lrp --semisup_dataset --add_classification --iterative_gradients --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 15000 --num_labels 20 --uniform_masks 1.0 --save_folder lrp0_resnet50_lab20_lr1e-5_s42
+python train.py --model unrolled_lrp --semisup_dataset --add_classification --iterative_gradients --batch_size 10 --pretrain_weight_name ./snapshot/resnet50_10_pre_train_21 --encoder resnet50 --seed 42 --epochs 100 --save_folder lrp0_resnet50_lab1464_lr1e-5_s42
 ```
 
-- TODO write conda environment for running the code.
+Note: The way for counting epochs differs between the semisupervised and supervised datasets. Therefore, you need to set a higher value for the epochs argument when training the UNet.
 
 
 ## Citation
 
-- TODO
+Yu, Xiaoyan, et al. "Model Guidance via Explanations Turns Image Classifiers into Segmentation Models." World Conference on Explainable Artificial Intelligence. Cham: Springer Nature Switzerland, 2024.
 
 
 ## Contact

models/multilabel.py

@@ -46,7 +46,7 @@ def __init__(self, args, logger):
             self.train_sampler = semi_supervised_sampler(self.data.training_dataset.idx_mask, self.data.training_dataset.idx_no_mask,
                                                         args.batch_size, len(self.data.training_dataset),0.5, 1.0,args.seed)
             self.train_loader=torch.utils.data.DataLoader(self.data.training_dataset,sampler=self.train_sampler.sample(),batch_size=args.batch_size)
-            self.test_loader=torch.utils.data.DataLoader(self.data.testing_dataset, batch_size=40)
+            self.test_loader=torch.utils.data.DataLoader(self.data.testing_dataset, batch_size=10)
 
 
         if args.pre_batch_size>0 and args.pre_epochs>0:
@@ -193,7 +193,7 @@ def pretraining_test(self,epoch):
 
         test_loss /= len(self.pre_test_loader.dataset)
 
-        _,_,_,_,match_ratio,avg_class_acc,_,_,_ = gen_log_message("pre-test",epoch,test_loss, metrics,len(self.pre_test_loader.dataset),return_value=True,print_class_iou=False)
+        _,_,match_ratio,avg_class_acc,_,_,_ = gen_log_message("pre-test",epoch,test_loss, metrics,len(self.pre_test_loader.dataset),return_value=True,print_class_iou=False)
 
 
         if self.args.wandb != 'None':
@@ -217,7 +217,7 @@ def train(self, epoch):
 
         for batch_idx, (data, class_labels,sem_gts,sem_gt_exist) in enumerate(self.train_loader): 
 
-            # measure number of samples before it goes to if function about self.args.only_send_labeled_data
+            # measure number of samples 
             num_samples=len(data)
             # 1) prepare data
             data, class_labels, sem_gts = data.to(self.device),class_labels.to(self.device), sem_gts.long().to(self.device)
@@ -239,10 +239,6 @@ def train(self, epoch):
                         loss1.backward(retain_graph=True) # backpropagate heatmap loss individually
                 loss2.backward() # backpropagate classification loss
             else:
-                if self.args.only_send_labeled_data: #only design for ablation test
-                    data=data[sem_gt_exist==True]
-                    class_labels=class_labels[sem_gt_exist==True]
-                    sem_gts=sem_gts[sem_gt_exist==True]
                 class_scores,heatmaps = self.model(data,class_labels)
                 loss, loss1, loss2= self.loss_function(class_scores,class_labels,heatmaps,sem_gts,sem_gt_exist)
 
@@ -285,7 +281,6 @@ def train(self, epoch):
     def test(self, epoch):
 
         self.model.eval()
-        # init parameters etc. for tracking
         test_loss = 0
 
         metrics=init_metrics(self.test_metrics,self.args.num_classes)
@@ -315,6 +310,9 @@ def test(self, epoch):
 
 
         test_loss /= len(self.test_loader.dataset)
+
+        # Reminder: match_ratio and avg_class_acc are not good for accessing the classification perfromance
+        # we check F1 score after saving the best checkpoint for segmentation task
         mIoU,iou_class,match_ratio,avg_class_acc,class_acc,pixel_acc,ap = gen_log_message("test",epoch,test_loss, metrics,len(self.test_loader.dataset),return_value=True,print_class_iou=True)