Skip to content

Commit

Permalink
Finalize imagenet
Browse files Browse the repository at this point in the history
  • Loading branch information
@pomonam
pomonam committed Jun 26, 2024
1 parent fb0ecfa commit f1e341f
Show file tree
Hide file tree
Showing 4 changed files with 57 additions and 38 deletions.
81 changes: 50 additions & 31 deletions examples/imagenet/README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@ To get started, please install the necessary packages by running:
pip install -r requirements.txt
```

## Training

We will use the pre-trained ResNet-50 model from `torchvision.models.resnet50`.

## Computing Pairwise Influence Scores
Expand All @@ -19,41 +17,64 @@ To compute pairwise influence scores on 1000 query data points using the `ekfac`
python analyze.py --dataset_dir PATH_TO_IMAGENET \
--query_gradient_rank -1 \
--query_batch_size 100 \
--train_batch_size 512 \
--train_batch_size 300 \
--factor_strategy ekfac
```

Replace `PATH_TO_IMAGENET` with the path to your ImageNet dataset directory.
On an A100 (80GB) GPU, it takes approximately 11 hours to compute the pairwise scores (including computing EKFAC factors):
Replace `PATH_TO_IMAGENET` with the path to your ImageNet dataset directory. On an A100 (80GB) GPU, it takes approximately 11 hours to compute the pairwise scores (including computing EKFAC factors):

```
----------------------------------------------------------------------------------------------------------------------------------
| Action | Mean duration (s) | Num calls | Total time (s) | Percentage % |
----------------------------------------------------------------------------------------------------------------------------------
| Total | - | 11 | 4.072e+04 | 100 % |
----------------------------------------------------------------------------------------------------------------------------------
| Compute Pairwise Score | 3.9136e+04 | 1 | 3.9136e+04 | 96.111 |
| Fit Lambda | 1056.7 | 1 | 1056.7 | 2.5952 |
| Fit Covariance | 454.05 | 1 | 454.05 | 1.1151 |
| Save Pairwise Score | 49.141 | 1 | 49.141 | 0.12068 |
| Perform Eigendecomposition | 10.589 | 1 | 10.589 | 0.026006 |
| Save Covariance | 5.7108 | 1 | 5.7108 | 0.014025 |
| Save Eigendecomposition | 5.5442 | 1 | 5.5442 | 0.013616 |
| Save Lambda | 0.97504 | 1 | 0.97504 | 0.0023945 |
| Load Eigendecomposition | 0.52786 | 1 | 0.52786 | 0.0012963 |
| Load Covariance | 0.30319 | 1 | 0.30319 | 0.00074458 |
| Load All Factors | 0.15671 | 1 | 0.15671 | 0.00038485 |
----------------------------------------------------------------------------------------------------------------------------------
```

Query batching (low-rank approximation to the query gradient; see **Section 3.2.2** from the paper) can be used to compute influence scores with a larger query batch size:

```bash
python analyze.py --dataset_dir PATH_TO_IMAGENET \
--query_gradient_rank 32 \
--query_batch_size 500 \
--train_batch_size 512 \
--query_batch_size 256 \
--train_batch_size 256 \
--factor_strategy ekfac
```

On an A100 (80GB) GPU, it takes roughly 3.5 hours to compute the pairwise scores with query batching (including computing EKFAC factors):

```
----------------------------------------------------------------------------------------------------------------------------------
| Action | Mean duration (s) | Num calls | Total time (s) | Percentage % |
----------------------------------------------------------------------------------------------------------------------------------
| Total | - | 3 | 9896.8 | 100 % |
----------------------------------------------------------------------------------------------------------------------------------
| Compute Pairwise Score | 9849.7 | 1 | 9849.7 | 99.524 |
| Save Pairwise Score | 47.075 | 1 | 47.075 | 0.47566 |
| Load All Factors | 0.014463 | 1 | 0.014463 | 0.00014614 |
----------------------------------------------------------------------------------------------------------------------------------
```

Assuming you ran the above two commands, `query_batching_analysis.py` contains code to compute the correlations between the full-rank and low-rank scores.

<p align="center">
<a href="#"><img width="380" img src="figure/query_batching.png" alt="Counterfactual"/></a>
<a href="#"><img width="380" img src="figure/query_batching.png" alt="Query Batching"/></a>
</p>

The averaged correlations between the low-rank and full rank scores for 100 data points is 0.95.
For more efficient computation, use half precision:
The averaged correlations between the low-rank and full rank scores for 100 data points is 0.94.
For even more efficient computation, use half precision:

```bash
python analyze.py --dataset_dir PATH_TO_IMAGENET \
Expand All @@ -64,41 +85,39 @@ python analyze.py --dataset_dir PATH_TO_IMAGENET \
--use_half_precision
```

This reduces computation time to about 20 minutes on an A100 (80GB) GPU:
This reduces computation time to about 85 minutes on an A100 (80GB) GPU:

```
----------------------------------------------------------------------------------------------------------------------------------
| Action | Mean duration (s) | Num calls | Total time (s) | Percentage % |
----------------------------------------------------------------------------------------------------------------------------------
| Total | - | 11 | 1211.8 | 100 % |
| Total | - | 11 | 4926.3 | 100 % |
----------------------------------------------------------------------------------------------------------------------------------
| Compute Pairwise Score | 1034.5 | 1 | 1034.5 | 85.368 |
| Fit Lambda | 88.231 | 1 | 88.231 | 7.2811 |
| Fit Covariance | 59.746 | 1 | 59.746 | 4.9305 |
| Perform Eigendecomposition | 14.831 | 1 | 14.831 | 1.2239 |
| Save Covariance | 5.8912 | 1 | 5.8912 | 0.48617 |
| Save Eigendecomposition | 5.7726 | 1 | 5.7726 | 0.47638 |
| Save Lambda | 1.624 | 1 | 1.624 | 0.13402 |
| Load Covariance | 0.34494 | 1 | 0.34494 | 0.028465 |
| Load Eigendecomposition | 0.33595 | 1 | 0.33595 | 0.027724 |
| Load All Factors | 0.26719 | 1 | 0.26719 | 0.022049 |
| Save Pairwise Score | 0.26006 | 1 | 0.26006 | 0.021461 |
| Compute Pairwise Score | 4446.0 | 1 | 4446.0 | 90.249 |
| Fit Lambda | 255.45 | 1 | 255.45 | 5.1853 |
| Fit Covariance | 186.86 | 1 | 186.86 | 3.7931 |
| Save Pairwise Score | 23.205 | 1 | 23.205 | 0.47104 |
| Perform Eigendecomposition | 7.1356 | 1 | 7.1356 | 0.14485 |
| Save Eigendecomposition | 3.3045 | 1 | 3.3045 | 0.067079 |
| Save Covariance | 2.993 | 1 | 2.993 | 0.060756 |
| Save Lambda | 0.58278 | 1 | 0.58278 | 0.01183 |
| Load Eigendecomposition | 0.39114 | 1 | 0.39114 | 0.0079398 |
| Load Covariance | 0.27701 | 1 | 0.27701 | 0.005623 |
| Load All Factors | 0.1699 | 1 | 0.1699 | 0.0034489 |
----------------------------------------------------------------------------------------------------------------------------------
```


## Computing Pairwise Influence Scores with DDP

You can also use [DistributedDataParallel (DDP)](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) to speed up influence computations. To run influence analysis with four A100 (80GB) GPUs and query batching, use the command:

```bash
torchrun --standalone --nnodes=1 --nproc-per-node=2 ddp_analyze.py --dataset_dir PATH_TO_IMAGENET \
torchrun --standalone --nnodes=1 --nproc-per-node=4 ddp_analyze.py --dataset_dir PATH_TO_IMAGENET \
--query_gradient_rank 32 \
--factor_batch_size 512 \
--query_batch_size 100 \
--train_batch_size 512 \
--factor_strategy ekfac \
--half_precision
--train_batch_size 256 \
--factor_strategy ekfac
```

It takes approximately 1 hour to compute the pairwise scores:
Expand All @@ -123,4 +142,4 @@ It takes approximately 1 hour to compute the pairwise scores:
----------------------------------------------------------------------------------------------------------------------------------
```

You can use more GPUs to further speed up the influence computations.
You can use more GPUs to further speed up the influence computations (or use AMP + half precision).
5 changes: 1 addition & 4 deletions examples/imagenet/analyze.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
import argparse
import logging
from typing import Tuple

import torch

Expand All @@ -12,8 +11,6 @@
from kronfluence.utils.common.score_arguments import all_low_precision_score_arguments
from kronfluence.utils.dataset import DataLoaderKwargs

BATCH_TYPE = Tuple[torch.Tensor, torch.Tensor]


def parse_args():
parser = argparse.ArgumentParser(description="Influence analysis on ImageNet dataset.")
Expand Down Expand Up @@ -46,7 +43,7 @@ def parse_args():
parser.add_argument(
"--train_batch_size",
type=int,
default=128,
default=256,
help="Batch size for computing training gradient.",
)
parser.add_argument(
Expand Down
1 change: 0 additions & 1 deletion examples/imagenet/ddp_analyze.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
from kronfluence.utils.model import apply_ddp

torch.backends.cudnn.benchmark = True
BATCH_DTYPE = Tuple[torch.Tensor, torch.Tensor]
LOCAL_RANK = int(os.environ["LOCAL_RANK"])
WORLD_RANK = int(os.environ["RANK"])
WORLD_SIZE = int(os.environ["WORLD_SIZE"])
Expand Down
8 changes: 6 additions & 2 deletions examples/imagenet/query_batching_analysis.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,12 @@ def main():
logging.basicConfig(level=logging.INFO)

# Load the scores. You might need to modify the path.
full_scores = Analyzer.load_file("scores_ekfac/pairwise_scores.safetensors")["all_modules"]
lr_scores = Analyzer.load_file("scores_ekfac_qlr32/pairwise_scores.safetensors")["all_modules"]
full_scores = Analyzer.load_file("influence_results/imagenet/scores_ekfac/pairwise_scores.safetensors")[
"all_modules"
]
lr_scores = Analyzer.load_file("influence_results/imagenet/scores_ekfac_qlr32/pairwise_scores.safetensors")[
"all_modules"
]

# Only plot first 1000 points to avoid clutter.
plt.rcParams.update({"figure.dpi": 150})
Expand Down

0 comments on commit f1e341f

Please sign in to comment.