About the Transfer entropy code

[guoronggeng]

hello! I can't find the code that caculate the transfer entropy of attention layer. Could you please tell me which file it's in?

[sihaoevery]

Hi,

Thanks for your attention to our work.

The code snippet you are looking for is at lines 323 and 335 of models_shrink.py.

1. Switch to True to turn on the entropy calculation.
2. Change the line 83 x = 2*x to x = x.
3. Manually set the dataloader to training dataloader when it runs inference in main_lambda.py.

You can refer to the test.sh, e.g., Model=deit_base_patch16_224_shrink Done=[] CKPT=deit_base_patch16_224-b5f2ef4d.pth python main_lambda.py \ --data-set IMNET \ --data-path path/to/imagenet/ \ --model $Model \ --prune_layer [] \ --done_layer $Done \ --seed 42 \ --num_workers 10 \ --batch-size 256 \ --finetune $CKPT \ --eval \ --no-model-ema

The script will calculate the entropy at the cls head and give the inference result at the training set. Similarly, if Done is set to [0,3], it will calculate the entropy when 0 and 3 attention layers are skipped. Then the subtraction is the transfer entropy of layers [0,3].

Please remember to turn back the line 83 if you'd like to evaluate our model.

Sihao

[sihaoevery]

We presume that testing data is not available. It's fine not to take step 3 (I know it's tricky!) if you just would like to calculate the transfer entropy.