From d7756a8f093e2aa1c3e66a46911d958775d59e34 Mon Sep 17 00:00:00 2001 From: Houman Safaai Date: Sat, 4 Oct 2025 11:24:04 -0400 Subject: [PATCH] Add GitHub CI/CD pipeline - 8 workflows: test, lint, security, docs, release, GPU, benchmark, deps - Pre-commit hooks and dependabot automation - Clean professional structure --- .github/PULL_REQUEST_TEMPLATE.md | 29 + .github/dependabot.yml | 27 + .github/workflows/benchmark.yml | 39 + .github/workflows/dependency-update.yml | 41 + .github/workflows/docs.yml | 86 +- .github/workflows/gpu-test.yml | 40 + .github/workflows/lint.yml | 44 + .github/workflows/pre-commit.yml | 27 + .github/workflows/release.yml | 51 + .github/workflows/security.yml | 45 + .github/workflows/test.yml | 53 + .gitignore | 139 +- .pre-commit-config.yaml | 48 + _arxiv/_archive/alignment_experiments.py.bak | 1288 +++++++++++++++++ .../alignment_preref/alignment_metrics.py | 243 ++++ .../alignment_preref/alignment_metrics_rem.py | 43 + _arxiv/_archive/alignment_preref/config.py | 116 ++ _arxiv/_archive/alignment_preref/datasets.py | 280 ++++ _arxiv/_archive/alignment_preref/plotting.py | 861 +++++++++++ .../_archive/alignment_preref/plotting_rem.py | 576 ++++++++ .../_archive/alignment_preref/processing.py | 869 +++++++++++ _arxiv/_archive/alignment_preref/train.py | 424 ++++++ _arxiv/_archive/alignment_preref/utils.py | 631 ++++++++ _arxiv/_archive/alignment_preref/utils_rem.py | 57 + _arxiv/_archive/alignment_stats.py | 97 ++ _arxiv/_archive/alignment_v2/__init__.py | 0 _arxiv/_archive/alignment_v2/datasets.py | 322 +++++ _arxiv/_archive/alignment_v2/files.py | 55 + _arxiv/_archive/alignment_v2/plotting.py | 565 ++++++++ _arxiv/_archive/alignment_v2/processing.py | 316 ++++ _arxiv/_archive/alignment_v2/train.py | 813 +++++++++++ _arxiv/_archive/alignment_v2/utils.py | 934 ++++++++++++ _arxiv/_archive/benchmark_dropout.py | 230 +++ _arxiv/_archive/config_alignment_stats.yaml | 56 + .../_archive/config_alignment_stats_old.yaml | 66 + .../_archive/config_alignment_stats_v2.yaml | 74 + _arxiv/_archive/config_pruning_modes.yaml | 80 + _arxiv/_archive/config_refactored_test.yaml | 65 + _arxiv/_archive/config_with_eigenvector.yaml | 59 + .../configs/config_alignment_experiment.yaml | 203 +++ .../configs/config_alignment_stats.yaml | 56 + .../configs/config_alignment_stats_old.yaml | 66 + .../configs/config_alignment_stats_v2.yaml | 74 + .../configs/config_pruning_modes.yaml | 80 + .../configs/config_refactored_test.yaml | 65 + .../configs/config_with_eigenvector.yaml | 59 + _arxiv/_archive/configs/training_example.yaml | 47 + _arxiv/_archive/debug_pruning_strategies.py | 235 +++ _arxiv/_archive/md_files/DOCUMENTATION.md | 52 + .../md_files/README_tensorized_dropout.md | 93 ++ .../_archive/md_files/REFACTORING_METRICS.md | 66 + _arxiv/archive/benchmarks/README.md | 29 + .../benchmark_dropout_strategies.py | 178 +++ .../benchmarks/benchmark_network_training.py | 239 +++ .../cluster/alignment-ddp-example.slurm | 157 ++ _arxiv/archive/cluster/alignment_stats.slurm | 42 + _arxiv/archive/cluster/ddp-example/README.md | 13 + _arxiv/archive/cluster/ddp-example/ddp.py | 218 +++ _arxiv/archive/cluster/ddp-example/ddp.slurm | 56 + .../cluster/imagenet-example/imagenet.py | 217 +++ .../cluster/imagenet-example/imagenet.slurm | 50 + .../lightning-example/lightning_train.sbatch | 35 + .../lightning-example/submit_lightning.sh | 26 + .../train_model_lightning.py | 110 ++ _arxiv/archive/cluster/slurm_settings.txt | 3 + .../cluster/testing-examples/test_job.slurm | 32 + .../refactoring_docs/GAUSSIAN_MI_SUMMARY.md | 84 ++ _arxiv/archive/scripts/README.md | 35 + _arxiv/archive/scripts/__init__.py | 0 _arxiv/archive/scripts/continual_mnist.py | 144 ++ _arxiv/archive/scripts/direct_pruning_test.py | 257 ++++ _arxiv/archive/scripts/run_benchmark.sh | 42 + _arxiv/archive/scripts/run_cascading_test.py | 73 + _arxiv/archive/scripts/run_cascading_test.sh | 50 + .../scripts/run_cascading_with_plots.py | 179 +++ .../archive/scripts/run_fixed_experiment.py | 402 +++++ .../scripts/run_multi_strategy_experiment.py | 614 ++++++++ _arxiv/archive/scripts/teacher_student.py | 93 ++ .../timetests/test_batched_alignment.py | 41 + .../scripts/timetests/test_dataloader.py | 106 ++ drafts/alignment_notes | 1 + requirements-dev.txt | 34 + requirements.in | 37 + 83 files changed, 14342 insertions(+), 140 deletions(-) create mode 100644 .github/PULL_REQUEST_TEMPLATE.md create mode 100644 .github/dependabot.yml create mode 100644 .github/workflows/benchmark.yml create mode 100644 .github/workflows/dependency-update.yml create mode 100644 .github/workflows/gpu-test.yml create mode 100644 .github/workflows/lint.yml create mode 100644 .github/workflows/pre-commit.yml create mode 100644 .github/workflows/release.yml create mode 100644 .github/workflows/security.yml create mode 100644 .github/workflows/test.yml create mode 100644 .pre-commit-config.yaml create mode 100644 _arxiv/_archive/alignment_experiments.py.bak create mode 100644 _arxiv/_archive/alignment_preref/alignment_metrics.py create mode 100644 _arxiv/_archive/alignment_preref/alignment_metrics_rem.py create mode 100644 _arxiv/_archive/alignment_preref/config.py create mode 100644 _arxiv/_archive/alignment_preref/datasets.py create mode 100644 _arxiv/_archive/alignment_preref/plotting.py create mode 100644 _arxiv/_archive/alignment_preref/plotting_rem.py create mode 100644 _arxiv/_archive/alignment_preref/processing.py create mode 100644 _arxiv/_archive/alignment_preref/train.py create mode 100644 _arxiv/_archive/alignment_preref/utils.py create mode 100644 _arxiv/_archive/alignment_preref/utils_rem.py create mode 100644 _arxiv/_archive/alignment_stats.py create mode 100644 _arxiv/_archive/alignment_v2/__init__.py create mode 100644 _arxiv/_archive/alignment_v2/datasets.py create mode 100644 _arxiv/_archive/alignment_v2/files.py create mode 100644 _arxiv/_archive/alignment_v2/plotting.py create mode 100644 _arxiv/_archive/alignment_v2/processing.py create mode 100644 _arxiv/_archive/alignment_v2/train.py create mode 100644 _arxiv/_archive/alignment_v2/utils.py create mode 100644 _arxiv/_archive/benchmark_dropout.py create mode 100644 _arxiv/_archive/config_alignment_stats.yaml create mode 100644 _arxiv/_archive/config_alignment_stats_old.yaml create mode 100644 _arxiv/_archive/config_alignment_stats_v2.yaml create mode 100644 _arxiv/_archive/config_pruning_modes.yaml create mode 100644 _arxiv/_archive/config_refactored_test.yaml create mode 100644 _arxiv/_archive/config_with_eigenvector.yaml create mode 100644 _arxiv/_archive/configs/config_alignment_experiment.yaml create mode 100644 _arxiv/_archive/configs/config_alignment_stats.yaml create mode 100644 _arxiv/_archive/configs/config_alignment_stats_old.yaml create mode 100644 _arxiv/_archive/configs/config_alignment_stats_v2.yaml create mode 100644 _arxiv/_archive/configs/config_pruning_modes.yaml create mode 100644 _arxiv/_archive/configs/config_refactored_test.yaml create mode 100644 _arxiv/_archive/configs/config_with_eigenvector.yaml create mode 100644 _arxiv/_archive/configs/training_example.yaml create mode 100644 _arxiv/_archive/debug_pruning_strategies.py create mode 100644 _arxiv/_archive/md_files/DOCUMENTATION.md create mode 100644 _arxiv/_archive/md_files/README_tensorized_dropout.md create mode 100644 _arxiv/_archive/md_files/REFACTORING_METRICS.md create mode 100644 _arxiv/archive/benchmarks/README.md create mode 100755 _arxiv/archive/benchmarks/benchmark_dropout_strategies.py create mode 100755 _arxiv/archive/benchmarks/benchmark_network_training.py create mode 100644 _arxiv/archive/cluster/alignment-ddp-example.slurm create mode 100644 _arxiv/archive/cluster/alignment_stats.slurm create mode 100644 _arxiv/archive/cluster/ddp-example/README.md create mode 100644 _arxiv/archive/cluster/ddp-example/ddp.py create mode 100644 _arxiv/archive/cluster/ddp-example/ddp.slurm create mode 100644 _arxiv/archive/cluster/imagenet-example/imagenet.py create mode 100644 _arxiv/archive/cluster/imagenet-example/imagenet.slurm create mode 100644 _arxiv/archive/cluster/lightning-example/lightning_train.sbatch create mode 100644 _arxiv/archive/cluster/lightning-example/submit_lightning.sh create mode 100644 _arxiv/archive/cluster/lightning-example/train_model_lightning.py create mode 100644 _arxiv/archive/cluster/slurm_settings.txt create mode 100644 _arxiv/archive/cluster/testing-examples/test_job.slurm create mode 100644 _arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md create mode 100644 _arxiv/archive/scripts/README.md create mode 100644 _arxiv/archive/scripts/__init__.py create mode 100644 _arxiv/archive/scripts/continual_mnist.py create mode 100644 _arxiv/archive/scripts/direct_pruning_test.py create mode 100755 _arxiv/archive/scripts/run_benchmark.sh create mode 100644 _arxiv/archive/scripts/run_cascading_test.py create mode 100755 _arxiv/archive/scripts/run_cascading_test.sh create mode 100755 _arxiv/archive/scripts/run_cascading_with_plots.py create mode 100644 _arxiv/archive/scripts/run_fixed_experiment.py create mode 100644 _arxiv/archive/scripts/run_multi_strategy_experiment.py create mode 100644 _arxiv/archive/scripts/teacher_student.py create mode 100644 _arxiv/archive/scripts/timetests/test_batched_alignment.py create mode 100644 _arxiv/archive/scripts/timetests/test_dataloader.py create mode 160000 drafts/alignment_notes create mode 100644 requirements-dev.txt create mode 100644 requirements.in diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 00000000..86884564 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,29 @@ +## Description +Brief description of the changes in this PR. + +## Type of Change +- [ ] Bug fix (non-breaking change which fixes an issue) +- [ ] New feature (non-breaking change which adds functionality) +- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) +- [ ] Documentation update +- [ ] Performance improvement +- [ ] Code refactoring + +## Testing +- [ ] Tests pass locally +- [ ] New tests added for new functionality +- [ ] All existing tests pass +- [ ] Manual testing completed + +## Checklist +- [ ] Code follows the project's coding standards +- [ ] Self-review of the code completed +- [ ] Code is properly commented +- [ ] Documentation updated (if applicable) +- [ ] No new warnings or errors introduced + +## Related Issues +Closes #(issue number) + +## Additional Notes +Any additional information that reviewers should know. diff --git a/.github/dependabot.yml b/.github/dependabot.yml new file mode 100644 index 00000000..f7f0b780 --- /dev/null +++ b/.github/dependabot.yml @@ -0,0 +1,27 @@ +version: 2 +updates: + - package-ecosystem: "pip" + directory: "/" + schedule: + interval: "weekly" + day: "monday" + time: "09:00" + open-pull-requests-limit: 10 + reviewers: [] + assignees: [] + commit-message: + prefix: "chore" + include: "scope" + + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" + day: "monday" + time: "09:00" + open-pull-requests-limit: 5 + reviewers: [] + assignees: [] + commit-message: + prefix: "chore" + include: "scope" diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml new file mode 100644 index 00000000..98e8f7e1 --- /dev/null +++ b/.github/workflows/benchmark.yml @@ -0,0 +1,39 @@ +name: Performance Benchmarks + +on: + push: + branches: [ main ] + schedule: + - cron: '0 0 * * 0' # Weekly on Sunday + workflow_dispatch: + +jobs: + benchmark: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pytest-benchmark memory-profiler + + - name: Run benchmarks + run: | + pytest tests/benchmarks/ -v --benchmark-only --benchmark-save=benchmark + python scripts/run_benchmarks.py + + - name: Upload benchmark results + uses: actions/upload-artifact@v3 + with: + name: benchmark-results + path: | + .benchmarks/ + benchmark_results/ diff --git a/.github/workflows/dependency-update.yml b/.github/workflows/dependency-update.yml new file mode 100644 index 00000000..1b189a4c --- /dev/null +++ b/.github/workflows/dependency-update.yml @@ -0,0 +1,41 @@ +name: Dependency Updates + +on: + schedule: + - cron: '0 0 * * 1' # Weekly on Monday + workflow_dispatch: + +jobs: + update-dependencies: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install pip-tools + + - name: Update requirements + run: | + pip-compile requirements.in --upgrade + pip-compile requirements-dev.in --upgrade + + - name: Create Pull Request + uses: peter-evans/create-pull-request@v5 + with: + token: ${{ secrets.GITHUB_TOKEN }} + commit-message: 'chore: update dependencies' + title: 'Automated dependency updates' + body: | + This PR contains automated dependency updates. + + Please review the changes and test before merging. + branch: dependency-updates + delete-branch: true diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index f39dc42c..3e3a221d 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -1,77 +1,41 @@ -name: Build and Deploy Documentation +name: Documentation on: push: - branches: [ develop ] + branches: [ main, develop ] pull_request: - branches: [ develop ] - workflow_dispatch: - -permissions: - contents: read - pages: write - id-token: write - -# Use a unique concurrency group to avoid conflicts -concurrency: - group: "sphinx-docs-${{ github.workflow }}-${{ github.ref }}" - cancel-in-progress: true + branches: [ main, develop ] jobs: - build-docs: - name: Build Sphinx Documentation + docs: runs-on: ubuntu-latest + steps: - - name: Checkout code - uses: actions/checkout@v4 - + - uses: actions/checkout@v4 + - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v4 with: - python-version: '3.9' - + python-version: '3.10' + - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -e .[docs] - # Also install the main package dependencies - pip install -e . - + pip install -r requirements.txt + pip install sphinx sphinx-rtd-theme sphinx-autodoc-typehints + - name: Build documentation run: | - echo "Starting documentation build..." - cd docs - echo "Current directory: $(pwd)" - echo "Contents of source directory:" - ls -la source/ - make clean - echo "Building Sphinx documentation..." - sphinx-build -b html source build/html - echo "Creating .nojekyll file..." - touch build/html/.nojekyll - echo "Documentation build complete!" - echo "Contents of build directory:" - ls -la build/html/ - - - name: Setup Pages - if: github.event_name == 'push' && github.ref == 'refs/heads/develop' - uses: actions/configure-pages@v5 - - - name: Upload artifact - if: github.event_name == 'push' && github.ref == 'refs/heads/develop' - uses: actions/upload-pages-artifact@v3 + cd docs/ + sphinx-build -b html . _build/html + + - name: Check for broken links + run: | + pip install linkchecker + linkchecker docs/_build/html/ --check-extern + + - name: Upload documentation + uses: actions/upload-artifact@v3 with: - path: './docs/build/html' - - deploy-docs: - name: Deploy to GitHub Pages - if: github.event_name == 'push' && github.ref == 'refs/heads/develop' - runs-on: ubuntu-latest - needs: build-docs - environment: - name: github-pages - url: ${{ steps.deployment.outputs.page_url }} - steps: - - name: Deploy to GitHub Pages - id: deployment - uses: actions/deploy-pages@v4 \ No newline at end of file + name: documentation + path: docs/_build/html/ \ No newline at end of file diff --git a/.github/workflows/gpu-test.yml b/.github/workflows/gpu-test.yml new file mode 100644 index 00000000..d299b686 --- /dev/null +++ b/.github/workflows/gpu-test.yml @@ -0,0 +1,40 @@ +name: GPU Tests + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + workflow_dispatch: + +jobs: + gpu-test: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install CUDA dependencies + run: | + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb + sudo dpkg -i cuda-keyring_1.0-1_all.deb + sudo apt-get update + sudo apt-get -y install cuda-toolkit-12-0 + + - name: Install Python dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pytest pytest-cov + + - name: Run GPU tests + run: | + python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')" + pytest tests/test_gpu/ -v --tb=short + env: + CUDA_VISIBLE_DEVICES: 0 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 00000000..0ecd1060 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,44 @@ +name: Code Quality + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + +jobs: + lint: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install black isort flake8 mypy pre-commit + + - name: Run black (code formatting) + run: black --check --diff src/ tests/ scripts/ + + - name: Run isort (import sorting) + run: isort --check-only --diff src/ tests/ scripts/ + + - name: Run flake8 (linting) + run: flake8 src/ tests/ scripts/ --max-line-length=100 --ignore=E203,W503 + + - name: Run mypy (type checking) + run: mypy src/alignment/ --ignore-missing-imports --no-strict-optional + + - name: Check for TODO/FIXME comments + run: | + if grep -r "TODO\|FIXME" src/ tests/; then + echo "Found TODO/FIXME comments. Please address them." + exit 1 + fi diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml new file mode 100644 index 00000000..3cfacdb3 --- /dev/null +++ b/.github/workflows/pre-commit.yml @@ -0,0 +1,27 @@ +name: Pre-commit + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + +jobs: + pre-commit: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install pre-commit + run: | + python -m pip install --upgrade pip + pip install pre-commit + + - name: Run pre-commit + uses: pre-commit/action@v3.0.0 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 00000000..d6252330 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,51 @@ +name: Release + +on: + push: + tags: + - 'v*' + +jobs: + release: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install build twine + + - name: Build package + run: python -m build + + - name: Check package + run: twine check dist/* + + - name: Upload to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} + run: twine upload dist/* + + - name: Create GitHub Release + uses: actions/create-release@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + tag_name: ${{ github.ref }} + release_name: Release ${{ github.ref }} + body: | + ## Installation + + ```bash + pip install alignment-framework + ``` + draft: false + prerelease: false diff --git a/.github/workflows/security.yml b/.github/workflows/security.yml new file mode 100644 index 00000000..5086398a --- /dev/null +++ b/.github/workflows/security.yml @@ -0,0 +1,45 @@ +name: Security Scan + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + schedule: + - cron: '0 2 * * 1' # Weekly on Monday at 2 AM + +jobs: + security: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.10' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install bandit safety semgrep + + - name: Run bandit (security linting) + run: bandit -r src/ -f json -o bandit-report.json || true + + - name: Run safety (dependency vulnerability check) + run: safety check --json --output safety-report.json || true + + - name: Run semgrep (static analysis) + run: semgrep --config=auto src/ --json --output=semgrep-report.json || true + + - name: Upload security reports + uses: actions/upload-artifact@v3 + with: + name: security-reports + path: | + bandit-report.json + safety-report.json + semgrep-report.json diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 00000000..4e196803 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,53 @@ +name: Tests + +on: + push: + branches: [ main, develop ] + pull_request: + branches: [ main, develop ] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: [3.8, 3.9, 3.10, 3.11] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Cache pip dependencies + uses: actions/cache@v3 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }} + restore-keys: | + ${{ runner.os }}-pip- + + - name: Install system dependencies + run: | + sudo apt-get update + sudo apt-get install -y libgl1-mesa-glx libglib2.0-0 + + - name: Install Python dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pytest pytest-cov pytest-xdist + + - name: Run tests + run: | + pytest tests/ -v --cov=src/alignment --cov-report=xml --cov-report=html -n auto + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v3 + with: + file: ./coverage.xml + flags: unittests + name: codecov-umbrella + fail_ci_if_error: false diff --git a/.gitignore b/.gitignore index d3db0f63..c2cd311d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,39 +1,11 @@ -# Saved files -*.pkl -slurm-*.out -*.pt -*.mat - -# cluster jobs -jobs/ -# VSCODE -.vscode/ -data/ -logs/ -# WANDB -wandb/ -results/ -# Old refactored versions (no longer needed) -# src/alignment_preref/ -# src/alignment_v2/ -# src/alignment_refactor_v1/ -#notebooks/ -#notebooks/data/* -examples/data/ -examples/results/ -#notebooks/.ipynb_checkpoints/* -checkpoints/ -logs/ -notebooks/ # Byte-compiled / optimized / DLL files __pycache__/ -**/__pycache__/ *.py[cod] *$py.class -src/alignment.egg-info/ + # C extensions *.so -_arxiv/ + # Distribution / packaging .Python build/ @@ -48,6 +20,7 @@ parts/ sdist/ var/ wheels/ +pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg @@ -55,8 +28,6 @@ share/python-wheels/ MANIFEST # PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec @@ -77,7 +48,6 @@ coverage.xml *.py,cover .hypothesis/ .pytest_cache/ -cover/ # Translations *.mo @@ -100,7 +70,6 @@ instance/ docs/_build/ # PyBuilder -.pybuilder/ target/ # Jupyter Notebook @@ -109,28 +78,14 @@ target/ # IPython profile_default/ ipython_config.py -drafts/ # pyenv -# For a library or package, you might want to ignore these files since the code is -# intended to run in multiple environments; otherwise, check them in: -# .python-version +.python-version # pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# poetry -# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. -# This is especially recommended for binary packages to ensure reproducibility, and is more -# commonly ignored for libraries. -# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control -#poetry.lock - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow +Pipfile.lock + +# PEP 582 __pypackages__/ # Celery stuff @@ -167,33 +122,59 @@ dmypy.json # Pyre type checker .pyre/ -# pytype static type analyzer -.pytype/ - -# Cython debug symbols -cython_debug/ - -# PyCharm -# JetBrains specific template is maintainted in a separate JetBrains.gitignore that can -# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore -# and can be added to the global gitignore or merged into this file. For a more nuclear -# option (not recommended) you can uncomment the following to ignore the entire idea folder. -#.idea/ - -# Datasets -/data/ -*.tar.gz - -# Datasets -/data/ -*.tar.gz -*.zip -*.7z -*.npz +# Project specific +*.pkl +*.pickle +*.h5 +*.hdf5 *.pt *.pth -*.bin -results/** -checkpoints/** -logs/** +*.ckpt +*.safetensors + +# Data directories +data/ +datasets/ +models/ +checkpoints/ +logs/ +runs/ +outputs/ +results/ +experiments/ + +# Temporary files +*.tmp +*.temp +.DS_Store +Thumbs.db +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +.DS_Store? +._* +.Spotlight-V100 +.Trashes +ehthumbs.db +Thumbs.db + +# Security +.secrets.baseline +bandit-report.json +safety-report.json +semgrep-report.json + +# Benchmarks +.benchmarks/ +benchmark_results/ + +# Documentation +docs/_build/ +docs/build/ \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 00000000..4c32932b --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,48 @@ +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.4.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-added-large-files + - id: check-merge-conflict + - id: debug-statements + - id: check-docstring-first + + - repo: https://github.com/psf/black + rev: 23.7.0 + hooks: + - id: black + language_version: python3 + + - repo: https://github.com/pycqa/isort + rev: 5.12.0 + hooks: + - id: isort + args: ["--profile", "black"] + + - repo: https://github.com/pycqa/flake8 + rev: 6.0.0 + hooks: + - id: flake8 + args: [--max-line-length=100, --ignore=E203,W503] + + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.5.1 + hooks: + - id: mypy + additional_dependencies: [types-all] + args: [--ignore-missing-imports, --no-strict-optional] + + - repo: https://github.com/pycqa/bandit + rev: 1.7.5 + hooks: + - id: bandit + args: [-r, src/, -f, json, -o, bandit-report.json] + + - repo: https://github.com/Yelp/detect-secrets + rev: v1.4.0 + hooks: + - id: detect-secrets + args: ['--baseline', '.secrets.baseline'] diff --git a/_arxiv/_archive/alignment_experiments.py.bak b/_arxiv/_archive/alignment_experiments.py.bak new file mode 100644 index 00000000..0c510349 --- /dev/null +++ b/_arxiv/_archive/alignment_experiments.py.bak @@ -0,0 +1,1288 @@ +""" +Alignment experiment implementations. + +This module contains experiment classes for neural network alignment studies, +focusing on alignment metrics, dropout impacts, and training analysis. +""" + +import logging +import os +import sys +import argparse +from typing import Dict, List, Tuple, Optional, Any, Union +import copy +import pickle +import datetime +import json +import yaml + +import numpy as np +import torch +import torch.nn as nn +from tqdm import tqdm + +from alignment.config import ExperimentConfig +from alignment.experiments.experiment import Experiment +from alignment.metrics import AlignmentMetric, get_metric +from alignment.models.registry import create_model +from alignment.dropout import progressive_dropout, eigenvector_dropout +from alignment.training import train_model, evaluate_model +from alignment.utils.core import setup_logging +from alignment.utils.plotting import plot_dropout_results, plot_experiment_summary +from alignment.datasets import get_dataset, load_dataset + +logger = logging.getLogger(__name__) + + +class AlignmentExperiment(Experiment): + """ + Experiment class for studying neural network alignment properties. + + This class implements experiments that assess alignment between layers + in neural networks, with support for different dropout strategies, + multiple metrics, and visualization. + """ + + def __init__(self, config: Dict) -> None: + """Initialize the experiment with the given config. + + Args: + config: Experiment configuration. + """ + super().__init__(config) + + # Make sure dataset_config has required fields to prevent loading errors + if isinstance(self.config.dataset, dict): + dataset_config = self.config.dataset + + # Ensure consistent naming between dataset.name and dataset.dataset_name + if "name" in dataset_config and "dataset_name" not in dataset_config: + dataset_config["dataset_name"] = dataset_config["name"] + elif "dataset_name" in dataset_config and "name" not in dataset_config: + dataset_config["name"] = dataset_config["dataset_name"] + elif "name" not in dataset_config and "dataset_name" not in dataset_config: + if isinstance(self.config.dataset, str): + # If config.dataset is a string, use it as the dataset name + dataset_config["name"] = self.config.dataset + dataset_config["dataset_name"] = self.config.dataset + logger.info(f"Using dataset name from string: {self.config.dataset}") + else: + # Default to CIFAR10 if no dataset name provided + logger.warning("No dataset name provided in configuration. Using default: cifar10") + dataset_config["name"] = "cifar10" + dataset_config["dataset_name"] = "cifar10" + + logger.info(f"Using dataset: {dataset_config.get('name', 'unknown')}") + + # Set up experiment-specific paths + self.figure_path = None + self.weights_path = None + self.setup_paths() + + logger.debug(f"Initialized alignment experiment") + + # Add a sanity check + # Ensure checkpoint parameters exist in the config + if not hasattr(self.config, "checkpoint"): + self.config.checkpoint = {} + logger.warning("Checkpoint configuration not found, using defaults") + + self.metric = get_metric(config.alignment.metric) + + def get_basename(self) -> str: + """ + Get the base name for the experiment. + + Returns: + Base name string + """ + return f"alignment_{self.config.model.model_name}_{self.config.dataset.dataset_name}" + + def prepare_path(self) -> List[str]: + """ + Prepare the experiment path components. + + Returns: + List of path components + """ + return [ + "alignment", + self.config.model.model_name, + self.config.dataset.dataset_name, + f"metric_{self.config.alignment.metric}" + ] + + def create_networks(self) -> List[nn.Module]: + """ + Create multiple neural networks for the experiment, each with different initialization. + + Following the alignment_v2 approach, this creates multiple independent networks + rather than copies of a single network. + + Returns: + List containing multiple independently initialized networks + """ + # Get the number of network replicates from config + num_replicates = self.config.training.replicates if hasattr(self.config.training, "replicates") else 5 + + # Create multiple models based on configuration, each with different initialization + networks = [] + for i in range(num_replicates): + # Set a different seed for each network to ensure different initializations + if hasattr(self.config, 'seed') and self.config.seed is not None: + # Use the base seed plus the replicate index to get different but reproducible initializations + torch.manual_seed(self.config.seed + i) + torch.cuda.manual_seed_all(self.config.seed + i) + np.random.seed(self.config.seed + i) + + # Create a new model + model = create_model(self.config.model) + model.to(self.device) + networks.append(model) + + logger.info(f"Created {len(networks)} models with independent initializations: {self.config.model.model_name}") + + return networks + + def main(self) -> Tuple[Dict, List[nn.Module]]: + """ + Main experiment execution method. + + Returns: + Tuple of (results dictionary, list of networks) + """ + # Additional initialization + self.setup_paths() + logger.info(f"Set up paths. Results will be saved to {self.results_path}") + + # Create network models for experiment + networks = self.create_networks() + + # Run the experiment based on the specified type + experiment_type = getattr(self.config, 'experiment_type', 'alignment_analysis') + + if experiment_type == "alignment_analysis" or experiment_type == "alignment": + # Run alignment analysis and store results + results = self.run_alignment_analysis(networks) + + elif experiment_type == "progressive_dropout": + # Run progressive dropout experiment for all networks + results = self._run_progressive_dropout(networks) + + # Generate and log visualization + self._plot_dropout_results( + results, + self.figure_path, + title_prefix=f"{getattr(self.config, 'experiment_name', 'Progressive Dropout')}" + ) + + # Save results for later reference + with open(os.path.join(self.results_path, "progressive_dropout_results.pkl"), "wb") as f: + pickle.dump(results, f) + + elif experiment_type == "eigenvector_dropout": + # Run eigenvector dropout experiment for one network + results = self._run_eigenvector_dropout(networks[0]) + + # Generate and log visualization (this has a different format) + self._generate_dropout_plots( + results, + "Eigenvector Dropout", + "eigenvector_dropout" + ) + + # Save results for later reference + with open(os.path.join(self.results_path, "eigenvector_dropout_results.pkl"), "wb") as f: + pickle.dump(results, f) + + elif experiment_type == "training_alignment": + # Run training alignment experiment + # This is handled separately because it needs different structure + results = self._run_training_alignment() + + # Generate and log visualization + self._generate_training_alignment_plots(results) + + # Save results for later reference + with open(os.path.join(self.results_path, "training_alignment_results.pkl"), "wb") as f: + pickle.dump(results, f) + + else: + raise ValueError(f"Unsupported experiment type: {experiment_type}") + + logger.info(f"Completed {getattr(self.config, 'experiment_name', experiment_type)} experiment") + + return results, networks + + def _run_progressive_dropout(self, networks: List[nn.Module]) -> Dict: + """ + Run progressive dropout experiment with multiple networks. + + Uses the progressive_dropout function from alignment.dropout module + instead of reimplementing the logic directly. + + Args: + networks: List of networks to run the experiment on. + + Returns: + Dictionary with results that match the expected format for plotting. + """ + # Get dropout fractions from config + dropin_min = self.config.alignment.dropout_min + dropin_max = self.config.alignment.dropout_max + num_dropout_fractions = self.config.alignment.dropout_steps + dropout_fractions = np.linspace(dropin_min, dropin_max, num_dropout_fractions).tolist() + + # Get pruning mode and dropout mode from config + pruning_mode = getattr(self.config.extra, "dropout_pruning_mode", "global_joint") + dropout_mode = getattr(self.config.extra, "dropout_mode", "scaled") + + # Print exactly what mode is being used + logger.info(f"Running progressive dropout with pruning_mode={pruning_mode}, dropout_mode={dropout_mode}") + + # Debug metric type + logger.info(f"Using metric type: {type(self.metric)}") + + # Prepare dataset + try: + batch_size = getattr(self.config.dataset, "batch_size", 128) + dataset = load_dataset(self.config.dataset, batch_size=batch_size) + except Exception as e: + logger.error(f"Error loading dataset: {str(e)}") + return {"error": f"Dataset configuration error: {str(e)}", "dropout_fractions": dropout_fractions} + + # IMPORTANT: Train all networks at once before applying dropout + logger.info(f"Training {len(networks)} networks before applying progressive dropout") + + # Track training history for plotting + training_history = { + 'train_loss': [], + 'train_acc': [], + 'test_loss': [], + 'test_acc': [] + } + + # Get training parameters from config + num_epochs = getattr(self.config.training, "epochs", 5) + learning_rate = getattr(self.config.training, "learning_rate", 0.001) + + # Set up optimizer for each network + optimizers = [] + for network in networks: + network.to(self.device) + optimizers.append(torch.optim.Adam(network.parameters(), lr=learning_rate)) + + # Train all networks for each epoch - with progress bars + from tqdm import tqdm + epoch_pbar = tqdm(range(num_epochs), desc="Training epochs", position=0) + for epoch in epoch_pbar: + # Initialize epoch stats + epoch_train_loss = 0.0 + epoch_train_acc = 0.0 + epoch_test_loss = 0.0 + epoch_test_acc = 0.0 + + # Training phase + net_pbar = tqdm(enumerate(zip(networks, optimizers)), + desc=f"Epoch {epoch+1}/{num_epochs} networks", + total=len(networks), + position=1, + leave=False) + + for network_idx, (network, optimizer) in net_pbar: + network.train() + running_loss = 0.0 + correct = 0 + total = 0 + + # Train on each batch + batch_pbar = tqdm(dataset.train_loader, + desc=f"Network {network_idx+1}/{len(networks)}", + position=2, + leave=False) + + for inputs, targets in batch_pbar: + inputs, targets = inputs.to(self.device), targets.to(self.device) + + # Zero the parameter gradients + optimizer.zero_grad() + + # Forward pass + outputs = network(inputs) + loss = torch.nn.functional.cross_entropy(outputs, targets) + + # Backward pass and optimize + loss.backward() + optimizer.step() + + # Track statistics + running_loss += loss.item() + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + # Update batch progress bar + if total > 0: + batch_pbar.set_postfix({ + 'loss': running_loss / len(batch_pbar), + 'acc': 100.0 * correct / total + }) + + # Calculate network training statistics + if total > 0: + network_train_loss = running_loss / len(dataset.train_loader) + network_train_acc = 100.0 * correct / total + + # Accumulate for epoch average + epoch_train_loss += network_train_loss + epoch_train_acc += network_train_acc + + # Update network progress bar + net_pbar.set_postfix({ + 'train_loss': network_train_loss, + 'train_acc': network_train_acc + }) + + # Evaluation phase + network.eval() + test_correct = 0 + test_total = 0 + test_loss_sum = 0.0 + + # Evaluation progress bar + eval_pbar = tqdm(dataset.test_loader, + desc=f"Evaluating network {network_idx+1}", + position=2, + leave=False) + + with torch.no_grad(): + for inputs, targets in eval_pbar: + inputs, targets = inputs.to(self.device), targets.to(self.device) + outputs = network(inputs) + + # Calculate loss + loss = torch.nn.functional.cross_entropy(outputs, targets, reduction='sum') + test_loss_sum += loss.item() + + # Calculate accuracy + _, predicted = outputs.max(1) + test_total += targets.size(0) + test_correct += predicted.eq(targets).sum().item() + + # Update eval progress bar + if test_total > 0: + eval_pbar.set_postfix({ + 'test_loss': test_loss_sum / test_total, + 'test_acc': 100.0 * test_correct / test_total + }) + + # Calculate network testing statistics + if test_total > 0: + network_test_loss = test_loss_sum / test_total + network_test_acc = 100.0 * test_correct / test_total + + # Accumulate for epoch average + epoch_test_loss += network_test_loss + epoch_test_acc += network_test_acc + + # Calculate and store epoch averages + epoch_train_loss /= len(networks) + epoch_train_acc /= len(networks) + epoch_test_loss /= len(networks) + epoch_test_acc /= len(networks) + + training_history['train_loss'].append(epoch_train_loss) + training_history['train_acc'].append(epoch_train_acc) + training_history['test_loss'].append(epoch_test_loss) + training_history['test_acc'].append(epoch_test_acc) + + # Update epoch progress bar + epoch_pbar.set_postfix({ + 'train_loss': epoch_train_loss, + 'train_acc': f"{epoch_train_acc:.2f}%", + 'test_loss': epoch_test_loss, + 'test_acc': f"{epoch_test_acc:.2f}%" + }) + + # Log progress + logger.info(f"Epoch {epoch+1}/{num_epochs}: " + f"Train Loss={epoch_train_loss:.4f}, Train Acc={epoch_train_acc:.2f}%, " + f"Test Loss={epoch_test_loss:.4f}, Test Acc={epoch_test_acc:.2f}%") + + logger.info(f"Completed training {len(networks)} networks. Proceeding to progressive dropout...") + + # Initialize results structure with correct format for plotting + final_results = { + "dropout_fractions": dropout_fractions, + "accuracies": {"high_rq": [], "low_rq": [], "random": []}, + "losses": {"high_rq": [], "low_rq": [], "random": []}, + "stds": {"high_rq": [], "low_rq": [], "random": []}, + "training_history": training_history + } + + # Call the progressive_dropout function with our trained networks + from alignment.dropout import progressive_dropout + + try: + # Define the strategies to run + strategies = ["high_rq", "low_rq", "random"] + + # Run each strategy separately with progress bars + strategy_pbar = tqdm(strategies, desc="Running pruning strategies", position=0) + for strategy in strategy_pbar: + strategy_pbar.set_description(f"Strategy: {strategy}") + + # Run progressive dropout with this strategy + network_accuracies, network_losses = progressive_dropout( + networks, + dataset, + dropout_fractions, + self.metric, + self.device, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + strategy=strategy + ) + + # Calculate averages and standard deviations across networks for this strategy + fraction_pbar = tqdm(range(len(dropout_fractions)), + desc="Processing dropout fractions", + position=1, + leave=False) + + for fraction_idx in fraction_pbar: + fraction_pbar.set_description(f"Fraction: {dropout_fractions[fraction_idx]:.2f}") + + # Collect results for this fraction across all networks + fraction_accs = [] + fraction_losses = [] + + for net_idx in network_accuracies: + if fraction_idx < len(network_accuracies[net_idx]): + fraction_accs.append(network_accuracies[net_idx][fraction_idx]) + if fraction_idx < len(network_losses[net_idx]): + fraction_losses.append(network_losses[net_idx][fraction_idx]) + + # Calculate statistics + if fraction_accs: + mean_acc = np.mean(fraction_accs) + std_acc = np.std(fraction_accs) + mean_loss = np.mean(fraction_losses) if fraction_losses else 0.0 + + # Add to appropriate strategy in results + if fraction_idx >= len(final_results["accuracies"][strategy]): + final_results["accuracies"][strategy].append(mean_acc) + final_results["stds"][strategy].append(std_acc) + final_results["losses"][strategy].append(mean_loss) + else: + # Update existing value (shouldn't happen with our current code) + final_results["accuracies"][strategy][fraction_idx] = mean_acc + final_results["stds"][strategy][fraction_idx] = std_acc + final_results["losses"][strategy][fraction_idx] = mean_loss + + # Update progress bar + fraction_pbar.set_postfix({ + 'acc': f"{mean_acc:.2f}%", + 'std': f"{std_acc:.2f}%" + }) + + # Log results for verification + logger.info(f"Final results:") + for strategy in strategies: + logger.info(f" {strategy}: {final_results['accuracies'][strategy]}") + + except Exception as e: + logger.error(f"Error running progressive dropout: {str(e)}") + import traceback + logger.error(traceback.format_exc()) + return { + "error": f"Error running progressive dropout: {str(e)}", + "dropout_fractions": dropout_fractions, + "training_history": training_history + } + + # Generate dropout plots + try: + logger.info("Generating dropout plots with error bars") + from alignment.utils.plotting import plot_dropout_results + + saved_plots = plot_dropout_results( + final_results, + save_dir=self.figure_path, + title_prefix="Progressive Dropout", + pruning_mode=pruning_mode, + dropout_mode=dropout_mode + ) + + if saved_plots: + logger.info(f"Generated {len(saved_plots)} plots: {saved_plots}") + final_results["plot_files"] = saved_plots + else: + logger.warning("No plots were generated. Check your plot_dropout_results function.") + + except Exception as e: + logger.error(f"Error generating dropout plots: {str(e)}") + import traceback + logger.error(traceback.format_exc()) + + return final_results + + def _run_eigenvector_dropout(self, net: nn.Module) -> Dict: + """ + Run eigenvector dropout experiments. + + Uses the eigenvector_dropout function from alignment.dropout module + instead of reimplementing the logic directly. + + Args: + net: Neural network to analyze + + Returns: + Dictionary of eigenvector dropout results + """ + # Get dropout fractions from config + dropout_fractions = np.linspace( + self.config.alignment.dropout_min, + self.config.alignment.dropout_max, + self.config.alignment.dropout_steps + ).tolist() + + # Initialize result structure + results = { + "dropout_fractions": dropout_fractions, + "accuracies": {"eigenvector": []}, + "losses": {"eigenvector": []}, + "alignment_values": {"eigenvector": []} + } + + # Prepare dataset config + dataset_config = self.config.dataset + + # Process each dropout fraction + for dropout_fraction in tqdm(dropout_fractions, desc="Eigenvector Dropout"): + try: + # Call the eigenvector_dropout function directly + accuracy, alignment_values = eigenvector_dropout( + net, + dataset_config, + dropout_fraction=dropout_fraction, + metric=self.metric, + device=self.device, + dropout_mode=getattr(self.config.extra, "dropout_mode", "scaled"), + dropout_pruning_mode=getattr(self.config.extra, "dropout_pruning_mode", "global_joint") + ) + + # Store results + results["accuracies"]["eigenvector"].append(accuracy) + results["losses"]["eigenvector"].append(100.0 - accuracy) + results["alignment_values"]["eigenvector"].append(alignment_values) + + except Exception as e: + logger.error(f"Error in eigenvector dropout at fraction {dropout_fraction}: {str(e)}") + # Add placeholder values to maintain result structure + results["accuracies"]["eigenvector"].append(0.0) + results["losses"]["eigenvector"].append(100.0) + results["alignment_values"]["eigenvector"].append(None) + + return results + + def plot(self, results: Dict) -> Dict[str, List[str]]: + """ + Plot experiment results using enhanced visualization methods. + + Args: + results: Dictionary containing experiment results + + Returns: + Dictionary of plot paths by type + """ + logger.info("Generating experiment visualizations") + + # Check for results to plot + if not results: + logger.warning("No results to plot") + return {} + + plot_paths = {} + + # Generate experiment summary + try: + logger.info("Creating experiment summary visualization") + summary_path = self._generate_summary_plots(results) + if summary_path: + plot_paths["summary"] = summary_path + logger.info("Successfully created experiment summary visualization") + except Exception as e: + logger.error(f"Error creating summary visualization: {str(e)}", exc_info=True) + + # Generate progressive dropout plots + if "progressive_dropout" in results: + try: + logger.info("Generating progressive dropout visualizations") + prog_paths = self._plot_dropout_results( + results["progressive_dropout"], + self.figure_path, + "Progressive Dropout" + ) + if prog_paths: + plot_paths["progressive_dropout"] = prog_paths + logger.info("Successfully generated progressive dropout visualizations") + except Exception as e: + logger.error(f"Error generating progressive dropout visualizations: {str(e)}", exc_info=True) + + # Generate eigenvector dropout plots + if "eigenvector_dropout" in results: + try: + logger.info("Generating eigenvector dropout visualizations") + eig_paths = self._plot_dropout_results( + results["eigenvector_dropout"], + self.figure_path, + "Eigenvector Dropout" + ) + if eig_paths: + plot_paths["eigenvector_dropout"] = eig_paths + logger.info("Successfully generated eigenvector dropout visualizations") + except Exception as e: + logger.error(f"Error generating eigenvector dropout visualizations: {str(e)}", exc_info=True) + + # Generate combined visualizations if both types of results exist + if "progressive_dropout" in results and "eigenvector_dropout" in results: + try: + logger.info("Generating combined dropout comparison visualization") + comparison_path = self._generate_dropout_comparison( + results["progressive_dropout"], + results["eigenvector_dropout"] + ) + if comparison_path: + plot_paths["comparison"] = comparison_path + logger.info("Successfully generated combined dropout comparison") + except Exception as e: + logger.error(f"Error generating dropout comparison: {str(e)}", exc_info=True) + + logger.info(f"Completed experiment visualization generation: {len(plot_paths)} plot types") + return plot_paths + + def _plot_dropout_results(self, results, figure_path=None, title_prefix="Dropout"): + """ + Plot dropout experiment results using the plotting module. + + Args: + results: Results dictionary from progressive_dropout + figure_path: Directory to save plots to + title_prefix: Prefix for plot titles + + Returns: + List of saved plot files + """ + # Get pruning mode and dropout mode + pruning_mode = getattr(self.config.extra, "dropout_pruning_mode", "global_joint") + dropout_mode = getattr(self.config.extra, "dropout_mode", "scaled") + + # Ensure figure path exists + if figure_path is None: + figure_path = self.figure_path + + # Add debug logs to understand the results + logger.info(f"Plotting dropout results with pruning_mode={pruning_mode}, dropout_mode={dropout_mode}") + logger.info(f"Results keys: {list(results.keys())}") + logger.info(f"Figure path: {figure_path}") + + if "accuracies" in results: + logger.info(f"Accuracies keys: {list(results['accuracies'].keys())}") + for key in results["accuracies"]: + logger.info(f"Strategy {key} has {len(results['accuracies'][key])} data points") + if results["accuracies"][key]: + logger.info(f"First few values: {results['accuracies'][key][:3]}") + else: + logger.info(f"No data for {key}") + + # Generate plots using our custom function + try: + from alignment.utils.plotting import custom_plot_dropout + + # Make a deep copy of results to prevent modification + import copy + plot_results = copy.deepcopy(results) + + # Ensure we have data for at least one strategy (preferably high_rq) + if not plot_results.get("accuracies", {}).get("high_rq", []): + # Try to use data from any strategy that has data + for strategy in ["low_rq", "random", "eigenvector"]: + if plot_results.get("accuracies", {}).get(strategy, []): + # Use this data for high_rq + logger.info(f"Using {strategy} data for high_rq") + plot_results["accuracies"]["high_rq"] = plot_results["accuracies"][strategy] + + # Also copy to losses if needed + if strategy in plot_results.get("losses", {}): + plot_results["losses"]["high_rq"] = plot_results["losses"][strategy] + break + + # Try the numeric indices if we still don't have data + if not plot_results.get("accuracies", {}).get("high_rq", []) and hasattr(self, "results") and hasattr(self.results, "network_accuracies"): + # Look for numeric indices + for idx in range(3): # Try indices 0, 1, 2 + if idx in self.results.network_accuracies and self.results.network_accuracies[idx]: + strategy_name = {0: "high_rq", 1: "low_rq", 2: "random"}.get(idx, "high_rq") + logger.info(f"Using numeric index {idx} data for {strategy_name}") + + # Make sure the strategy key exists + if "accuracies" not in plot_results: + plot_results["accuracies"] = {} + + # Copy the data + plot_results["accuracies"][strategy_name] = self.results.network_accuracies[idx] + + # Print final plot data + logger.info(f"Passing the following to custom_plot_dropout:") + logger.info(f"- dropout_fractions: {len(plot_results.get('dropout_fractions', []))} points") + logger.info(f"- high_rq data: {len(plot_results.get('accuracies', {}).get('high_rq', []))} points") + + # Call the custom plotting function + saved_figures = custom_plot_dropout( + plot_results, + figure_path, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + title_prefix=title_prefix + ) + + # Check if any plots were generated + if saved_figures: + logger.info(f"Generated {len(saved_figures)} plots: {saved_figures}") + + # Log plots to wandb if configured + if hasattr(self.config.checkpointing, "use_wandb") and self.config.checkpointing.use_wandb: + try: + from alignment.utils.plotting import log_plots_to_wandb + log_plots_to_wandb(saved_figures) + logger.info(f"Logged {len(saved_figures)} dropout plots to wandb") + except Exception as e: + logger.warning(f"Failed to log dropout plots to wandb: {str(e)}") + else: + logger.warning("No plots were generated by custom_plot_dropout") + + return saved_figures + + except Exception as e: + logger.error(f"Error generating dropout plots: {str(e)}") + import traceback + logger.error(traceback.format_exc()) + return [] + + def _generate_dropout_plots(self, results, title, filename=None): + """ + Generate enhanced dropout plots. + + Args: + results: Dropout experiment results + title: Plot title + filename: Base filename (optional) + + Returns: + List of saved plot files + """ + # Call the _plot_dropout_results function with our figure path and title + return self._plot_dropout_results(results, self.figure_path, title) + + def _generate_summary_plots(self, results): + """ + Generate comprehensive summary plots for the experiment using the plotting module. + + Args: + results: Experiment results dictionary + + Returns: + Path to the saved summary plot + """ + try: + from alignment.utils.plotting import plot_experiment_summary + + # Generate the summary plot + experiment_name = getattr(self.config, "experiment_name", self.get_basename()) + filepath = plot_experiment_summary( + results, + self.figure_path, + experiment_name=experiment_name + ) + + # Log to wandb if configured + if filepath and hasattr(self.config.checkpointing, "use_wandb") and self.config.checkpointing.use_wandb: + try: + import wandb + if wandb.run is not None: + wandb.log({"experiment_summary": wandb.Image(filepath)}) + logger.info("Logged experiment summary to wandb") + except Exception as e: + logger.warning(f"Failed to log experiment summary to wandb: {str(e)}") + + return filepath + + except Exception as e: + logger.error(f"Error creating summary visualization: {str(e)}") + return None + + def _generate_dropout_comparison(self, progressive_results, eigenvector_results): + """ + Generate comparison plots between progressive and eigenvector dropout using the plotting module. + + Args: + progressive_results: Progressive dropout results + eigenvector_results: Eigenvector dropout results + + Returns: + Path to the saved comparison plot + """ + try: + from alignment.utils.plotting import plot_dropout_comparison + + # Generate the comparison plot + experiment_name = getattr(self.config, "experiment_name", "Experiment") + filepath = plot_dropout_comparison( + progressive_results, + eigenvector_results, + self.figure_path, + title=f"Dropout Comparison: {experiment_name}" + ) + + # Log to wandb if configured + if filepath and hasattr(self.config.checkpointing, "use_wandb") and self.config.checkpointing.use_wandb: + try: + import wandb + if wandb.run is not None: + wandb.log({"dropout_comparison": wandb.Image(filepath)}) + logger.info("Logged dropout comparison to wandb") + except Exception as e: + logger.warning(f"Failed to log dropout comparison to wandb: {str(e)}") + + return filepath + + except Exception as e: + logger.error(f"Error creating dropout comparison: {str(e)}") + return None + + def _initialize_wandb(self) -> Optional[Any]: + """ + Initialize wandb for experiment tracking. + + Returns: + wandb module if initialized successfully, None otherwise + """ + if not getattr(self.config.checkpointing, "use_wandb", False): + return None + + try: + import wandb + + # Get wandb configuration + entity = getattr(self.config, "wandb_entity", None) + project = getattr(self.config, "wandb_project", "neural_alignment") + + # Generate timestamp if it doesn't exist + if not hasattr(self, 'timestamp'): + self.timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + + # Create a unique run name + if hasattr(self.config, "experiment_name") and self.config.experiment_name: + run_name = f"{self.config.experiment_name}" + else: + run_name = f"{self.config.model.model_name}_{self.config.dataset.dataset_name}_{self.timestamp}" + + # Convert the config to a dictionary for wandb + config_dict = self.config.to_dict() if hasattr(self.config, "to_dict") else vars(self.config) + + # Initialize wandb run + wandb.init( + project=project, + entity=entity, + name=run_name, + config=config_dict, + reinit=True # Allow reinitializing in the same process + ) + + # Set up wandb to watch models if available + if hasattr(self, "networks") and self.networks: + # Watch the first network + wandb.watch(self.networks[0], log="all", log_freq=10) + + logger.info(f"Initialized Weights & Biases tracking for run '{run_name}'") + return wandb + except Exception as e: + logger.warning(f"Failed to initialize wandb: {str(e)}") + return None + + def _log_to_wandb(self, data: Dict) -> bool: + """ + Log data to wandb. + + Args: + data: Data to log to wandb (dict with keys corresponding to metric names) + + Returns: + True if logging was successful, False otherwise + """ + if not getattr(self.config.checkpointing, "use_wandb", False): + return False + + try: + import wandb + if wandb.run is None: + wandb_module = self._initialize_wandb() + if wandb_module is None: + return False + + # Log scalar metrics + metrics_to_log = {} + for k, v in data.items(): + # Skip non-scalar values + if isinstance(v, (int, float)): + metrics_to_log[k] = v + elif isinstance(v, (list, tuple)) and len(v) == 1: + # Single-element list or tuple with a scalar + if isinstance(v[0], (int, float)): + metrics_to_log[k] = v[0] + + # Log the metrics to wandb + if metrics_to_log: + wandb.log(metrics_to_log) + + # Check if data contains progressive dropout results that we should plot + if "progressive_dropout" in data and "plot_files" in data and getattr(self.config.extra, "log_images", True): + # Find all plot files and log them to wandb + for plot_file in data["plot_files"]: + if os.path.exists(plot_file): + plot_name = os.path.basename(plot_file).replace(".png", "") + wandb.log({plot_name: wandb.Image(plot_file)}) + + # If 'accuracies' exists in the progressive_dropout results, log them as a line plot + if "accuracies" in data["progressive_dropout"] and "dropout_fractions" in data["progressive_dropout"]: + # Create a table to log + accuracies = np.array(data["progressive_dropout"]["accuracies"]) + dropout_fractions = data["progressive_dropout"]["dropout_fractions"] + + # Calculate mean and std + mean_accuracies = np.mean(accuracies, axis=0) + std_accuracies = np.std(accuracies, axis=0) + + # Prepare data table for the custom plot + table_data = [[x, y, y_std] for x, y, y_std in zip(dropout_fractions, mean_accuracies, std_accuracies)] + + # Log the data as a custom plot + wandb.log({ + "dropout_curve": wandb.Table( + data=table_data, + columns=["dropout_fraction", "mean_accuracy", "std_accuracy"] + ) + }) + + # Create interactive plot + dropout_data = [] + for i, frac in enumerate(dropout_fractions): + dropout_data.append([frac, mean_accuracies[i], mean_accuracies[i]-std_accuracies[i], mean_accuracies[i]+std_accuracies[i]]) + + wandb.log({ + "dropout_plot_interactive": wandb.plot.line_series( + xs=[x[0] for x in dropout_data], + ys=[[x[1] for x in dropout_data], # mean + [x[2] for x in dropout_data], # mean-std + [x[3] for x in dropout_data]], # mean+std + keys=["Mean Accuracy", "Lower Bound", "Upper Bound"], + title=f"Progressive Dropout ({self.config.extra.dropout_pruning_mode})", + xname="Dropout Fraction", + yname="Accuracy" + ) + }) + + return True + except Exception as e: + logger.warning(f"Failed to log to wandb: {str(e)}") + return False + + def run(self) -> Tuple[Dict, List[nn.Module]]: + """ + Run the experiment with custom logging and wandb setup. + + Returns: + Tuple containing (results, networks) + """ + # Setup paths for experiment + self.setup_paths() + self.prepare_path() + + # Set up logging and suppress warnings + set_logging_level(logging.INFO) + + # Set random seed if configured + if hasattr(self.config, 'seed') and self.config.seed is not None: + torch.manual_seed(self.config.seed) + torch.cuda.manual_seed_all(self.config.seed) + np.random.seed(self.config.seed) + logger.info(f"Set random seed to {self.config.seed}") + + # Initialize wandb explicitly before starting experiment + if hasattr(self.config, 'checkpointing') and self.config.checkpointing.use_wandb: + self._initialize_wandb() + + # Run the main experiment + results, networks = self.main() + + # Store results as an attribute for saving + self.results = results + self.networks = networks + + # Ensure all plots are saved and logged + self.plot(results) + + # Save experiment state + self.save() + + # Close wandb run if it exists + try: + import wandb + if wandb.run is not None: + wandb.finish() + logger.info("Closed wandb run") + except Exception as e: + logger.warning(f"Error closing wandb run: {str(e)}") + + return results, networks + + def run_alignment_analysis(self, networks: List[nn.Module]) -> Dict: + """ + Run alignment analysis experiment. + + This experiment analyzes alignment metrics across networks and can run + both progressive dropout and eigenvector dropout experiments based on + configuration. + + Args: + networks: List of neural networks to analyze + + Returns: + Results dictionary + """ + logger.info("Running alignment analysis experiment") + + # Prepare the results structure + results = { + "progressive_dropout": {}, + "eigenvector_dropout": {}, + "config": self.config, + } + + # Progressive dropout experiment + if self.config.alignment.run_progressive: + logger.info(f"Running progressive dropout experiment (mode: {getattr(self.config.extra, 'dropout_pruning_mode', 'global_joint')})") + try: + prog_results = self._run_progressive_dropout(networks) + results["progressive_dropout"] = prog_results + + # Add debug logging to inspect results structure + if prog_results: + logger.info(f"Progressive dropout results keys: {list(prog_results.keys())}") + if 'accuracies' in prog_results: + logger.info(f"Progressive dropout accuracies keys: {list(prog_results['accuracies'].keys())}") + + # Check if there was an error in the dataset configuration + if "error" in prog_results: + logger.error(f"Progressive dropout experiment failed: {prog_results['error']}") + else: + # Generate plots right after getting results + try: + logger.info("Generating progressive dropout visualizations") + # Add debug print before and after plotting + figure_path = self._plot_dropout_results( + results["progressive_dropout"], + self.figure_path, + title_prefix="Progressive Dropout" + ) + logger.info(f"Progressive dropout plot generated at {figure_path}") + except Exception as e: + logger.error(f"Error generating progressive dropout visualizations: {str(e)}", exc_info=True) + except Exception as e: + logger.error(f"Error running progressive dropout experiment: {str(e)}") + results["progressive_dropout"] = {"error": str(e)} + + # Eigenvector dropout experiment + if self.config.alignment.run_eigenvector: + logger.info("Running eigenvector dropout experiment") + try: + # Note: For eigenvector dropout, we use just the first network + eig_results = self._run_eigenvector_dropout(networks[0]) + results["eigenvector_dropout"] = eig_results + + # Add debug logging + if eig_results: + logger.info(f"Eigenvector dropout results keys: {list(eig_results.keys())}") + if 'accuracies' in eig_results: + logger.info(f"Eigenvector dropout accuracies keys: {list(eig_results['accuracies'].keys())}") + + # Generate plots right after getting results + try: + logger.info("Generating eigenvector dropout visualizations") + # Use the specialized method for eigenvector dropout + self._generate_dropout_plots( + results["eigenvector_dropout"], + "Eigenvector Dropout", + "eigenvector_dropout" + ) + logger.info("Successfully generated eigenvector dropout visualizations") + except Exception as e: + logger.error(f"Error generating eigenvector dropout visualizations: {str(e)}", exc_info=True) + except Exception as e: + logger.error(f"Error running eigenvector dropout experiment: {str(e)}") + results["eigenvector_dropout"] = {"error": str(e)} + + # Generate and log summary plots for all experiments + try: + logger.info("Generating experiment summary") + self._generate_summary_plots(results) + logger.info("Successfully generated experiment summary") + except Exception as e: + logger.error(f"Error generating experiment summary: {str(e)}", exc_info=True) + + return results + + def setup_paths(self): + """ + Set up paths for experiment outputs. + + Creates necessary directories for figures, weights, and results. + """ + # Create timestamp subdirectory if needed + if getattr(self.config, "use_timestamp", True): + timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + base_dir = os.path.join("results", f"{self.get_basename()}_{timestamp}") + else: + base_dir = os.path.join("results", self.get_basename()) + + # Create main result directory + os.makedirs(base_dir, exist_ok=True) + self.results_path = base_dir + + # Create figure directory + self.figure_path = os.path.join(base_dir, "figures") + os.makedirs(self.figure_path, exist_ok=True) + + # Create weights directory for saved models + self.weights_path = os.path.join(base_dir, "weights") + os.makedirs(self.weights_path, exist_ok=True) + + # Set the device to use (default to CUDA if available) + if hasattr(self.config, "device") and self.config.device: + self.device = torch.device(self.config.device) + else: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + logger.info(f"Set up paths: results={self.results_path}, figures={self.figure_path}, device={self.device}") + + def save(self): + """ + Save experiment state including results and configuration. + + This method saves: + 1. Experiment results as pickle and JSON + 2. Configuration in YAML format + 3. Any model checkpoints if available + """ + try: + # Create a results directory if it doesn't exist + os.makedirs(self.results_path, exist_ok=True) + + # Save experiment results if available + if hasattr(self, 'results') and self.results: + results_file = os.path.join(self.results_path, "results.pkl") + with open(results_file, "wb") as f: + pickle.dump(self.results, f) + logger.info(f"Saved experiment results to {results_file}") + + # Also try to save as JSON for human readability + try: + # Convert tensors and other non-serializable objects + def clean_for_json(obj): + if isinstance(obj, (torch.Tensor, np.ndarray)): + return obj.tolist() if hasattr(obj, 'tolist') else str(obj) + elif isinstance(obj, dict): + return {k: clean_for_json(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [clean_for_json(item) for item in obj] + elif isinstance(obj, tuple): + return tuple(clean_for_json(item) for item in obj) + else: + # Return string representation for other types + return str(obj) if not isinstance(obj, (str, int, float, bool, type(None))) else obj + + json_results = clean_for_json(self.results) + json_file = os.path.join(self.results_path, "results.json") + with open(json_file, "w") as f: + json.dump(json_results, f, indent=2) + logger.info(f"Saved readable results to {json_file}") + except Exception as e: + logger.warning(f"Could not save results as JSON: {str(e)}") + + # Save configuration + config_file = os.path.join(self.results_path, "config.yaml") + # Check if we can convert config to dict + if hasattr(self.config, 'to_dict'): + config_dict = self.config.to_dict() + with open(config_file, "w") as f: + yaml.dump(config_dict, f, default_flow_style=False) + logger.info(f"Saved configuration to {config_file}") + + logger.info("Experiment state saved successfully") + + except Exception as e: + logger.error(f"Error saving experiment state: {str(e)}") + # Don't re-raise the exception, just log it + + +def set_logging_level(level=logging.INFO): + """ + Set logging level for all loggers in the alignment package. + + Args: + level: Logging level (default: logging.INFO) + """ + # Set root logger level + logging.getLogger().setLevel(level) + + # Set level for alignment package loggers + logging.getLogger('alignment').setLevel(level) + + # Reduce verbosity for other loggers + logging.getLogger('matplotlib').setLevel(logging.WARNING) + logging.getLogger('PIL').setLevel(logging.WARNING) + logging.getLogger('torch').setLevel(logging.WARNING) + logging.getLogger('wandb').setLevel(logging.WARNING) + + # Suppress common warnings + import warnings + warnings.filterwarnings("ignore", category=UserWarning, message=".*pruning fraction.*") + warnings.filterwarnings("ignore", category=UserWarning, message=".*Combined pruning across layers.*") + warnings.filterwarnings("ignore", category=UserWarning, message=".*dropout of the entire tensor.*") + +def cli_main(): + """Command-line interface for running alignment experiments.""" + parser = argparse.ArgumentParser(description="Neural network alignment experiment") + parser.add_argument("--config", type=str, required=True, help="Path to config YAML") + parser.add_argument("--quiet", action="store_true", help="Reduce logging output") + args = parser.parse_args() + + # Set logging level based on quiet flag + if args.quiet: + set_logging_level(logging.WARNING) + else: + set_logging_level(logging.INFO) + + # Load configuration + config = ExperimentConfig.load(args.config) + + # Initialize and run experiment + experiment = AlignmentExperiment(config) + results, networks = experiment.run() + + return results, networks + + +if __name__ == "__main__": + cli_main() \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/alignment_metrics.py b/_arxiv/_archive/alignment_preref/alignment_metrics.py new file mode 100644 index 00000000..e5ae6c91 --- /dev/null +++ b/_arxiv/_archive/alignment_preref/alignment_metrics.py @@ -0,0 +1,243 @@ +# -------------------------------------------- +# alignment_metrics.py +# -------------------------------------------- + +import torch +from alignment.utils import smart_pca, get_device + +class AlignmentMetrics: + """ + Provides static methods for various alignment metrics, + including 'delta_alignment'. + """ + + @staticmethod + def RQ(input_, weight_, scale_by_norm=False): + """ + Rayleigh Quotient alignment measure: + proportion of variance in `input_` explained by each row of `weight_`. + """ + return alignment(input_, weight_, method="alignment", relative=True, scale_by_norm=scale_by_norm) + + @staticmethod + def MI_0(input_, weight_, scale_by_norm=False): + """ + Placeholder for mutual information approach - version 0 + (currently reuses alignment(...) as a stand-in). + """ + return alignment(input_, weight_, method="alignment", relative=True, scale_by_norm=scale_by_norm) + + @staticmethod + def MI_1(input_, weight_): + """ + Placeholder for mutual information approach - version 1 + """ + return torch.tensor(0.0) + + @staticmethod + def delta_alignment(net, layer_idx, layer_input, scale_by_norm=False): + """ + alignment of (W_current - W_init) with the input's covariance. + This references net._init_weights for the initial weight. + If net._init_weights is absent, returns zeros. + """ + if not hasattr(net, "_init_weights"): + weight_diff = torch.zeros_like(net.get_alignment_weights()[layer_idx]) + else: + init_w = net._init_weights[layer_idx] + current_w = net.get_alignment_weights()[layer_idx] + weight_diff = current_w - init_w + + weight_diff = weight_diff.flatten(start_dim=1) + return alignment(layer_input, weight_diff, method="alignment", relative=True, scale_by_norm=scale_by_norm) + + @staticmethod + def measure(input_, weight_, method="RQ", scale_by_norm=False): + """ + For a single method and a single layer's (input_, weight_). + """ + if method == "RQ": + return AlignmentMetrics.RQ(input_, weight_, scale_by_norm=scale_by_norm) + elif method == "MI_0": + return AlignmentMetrics.MI_0(input_, weight_, scale_by_norm=scale_by_norm) + elif method == "MI_1": + return AlignmentMetrics.MI_1(input_, weight_) + else: + raise ValueError(f"Unknown alignment method {method}") + + @staticmethod + def patchwise_alignment(inp, w, method="RQ", weigh_by_var=True): + """ + 'Patchwise' alignment for CNN. + - inp shape: [B, F, patches], e.g. B=mini-batch, F=inC*kH*kW, patches=outH*outW + - w shape: [outC, F], the flattened filter weights. + - We compute alignment for each patch's (B,F) input. + - Weighted by patch-level variance across the batch dimension. + """ + B, F, P = inp.shape + var_patches = torch.var(inp, dim=0, keepdim=False) # => (F, P) + patchwise_var = var_patches.sum(dim=0) # => shape (P,) + + all_patch_vals = [] + all_patch_vars = [] + + for p in range(P): + patch_data = inp[:, :, p] # shape => [B, F] + cc = torch.cov(patch_data.T) # shape (F, F) + num_ = torch.sum(torch.matmul(w, cc) * w, dim=1) + denom_ = torch.sum(w*w, dim=1) + + patch_rq = num_ / denom_ + if method == "RQ": + patch_rq = patch_rq / torch.trace(cc) + patch_weight = patchwise_var[p] if weigh_by_var else 1.0 + + all_patch_vals.append(patch_rq * patch_weight) + all_patch_vars.append(patch_weight) + + total_weight = torch.stack(all_patch_vars).sum() + sum_rq = torch.stack(all_patch_vals, dim=0).sum(dim=0) # shape => (outC,) + + if total_weight > 0: + final_rq = sum_rq / total_weight + else: + final_rq = sum_rq * 0 + return final_rq # shape => (outC,) + + @staticmethod + def measure_methods(net, images, methods, precomputed=True, scale_by_norm=False): + """ + For each layer in 'net', produce a dict {method -> tensor}. + 1) get_layer_inputs + 2) preprocess => unfold or patchwise + 3) flatten weights + 4) measure alignment or patchwise alignment + + Parameters: + - scale_by_norm: if True, scales covariance matrices by their norm before computing RQ + """ + # 1) Get raw inputs + raw_inputs = net.get_layer_inputs(images, precomputed=precomputed) + # 2) Unfold or flatten them + preprocessed = net._preprocess_inputs(raw_inputs, compress_convolutional=True) + # 3) Flatten weights + layer_weights = net.get_alignment_weights(flatten=True) + + results_per_layer = [] + for layer_idx, (inp, wgt) in enumerate(zip(preprocessed, layer_weights)): + layer_dict = {} + for m in methods: + if m == "delta_alignment": + val = AlignmentMetrics.delta_alignment(net, layer_idx, inp, scale_by_norm=scale_by_norm) + else: + # if net says "patchwise" & input is 3D => do patchwise + if net.cnn_mode == "patchwise" and inp.ndim == 3: + val = AlignmentMetrics.patchwise_alignment(inp, wgt, method=m, weigh_by_var=True) + else: + # single-cov + val = AlignmentMetrics.measure(inp, wgt, method=m, scale_by_norm=scale_by_norm) + layer_dict[m] = val + results_per_layer.append(layer_dict) + + return results_per_layer + + @staticmethod + def compute_eigenvalues(x): + """ + Do a standard PCA with 'smart_pca' to get eigenvalues of x (batch, features). + """ + w, v = smart_pca(x.T, centered=True) + return w, v + + @staticmethod + def measure_expected_distribution(method, eigenvals, bins=50, num_tests=100): + """ + Return (counts, bin_edges) for a random alignment distribution of 'method'. + """ + if method == "RQ": + counts, edges, _ = expected_alignment_distribution( + eigenvals, + relative=True, + valid_rotation=False, + with_rotation=True, + bins=bins, + num_tests=num_tests + ) + return counts, edges + elif method in ["MI_0", "MI_1", "delta_alignment"]: + return None, None + else: + return None, None + +def alignment(input, weight, method="alignment", relative=True, scale_by_norm=False): + """ + measure alignment by computing RQ = (w^T C w) / (w^T w), + optionally normalized by trace(C). + + - input shape can be (N, F) for single-cov approach, + or (B, F, patches) if you do patchwise in patchwise_alignment. + - weight shape (outC, F). + - scale_by_norm: if True, scales the covariance matrix by its Frobenius norm + before computing RQ (similar to "similarity" in alignment_v2) + """ + if input.ndim == 2: + # standard single-cov approach + cc = torch.cov(input.T) + + # Scale covariance by its norm if requested (like "similarity" in v2) + if scale_by_norm: + cc_norm = torch.norm(cc) + if cc_norm > 0: + cc = cc / cc_norm + + numerator = torch.sum(torch.matmul(weight, cc) * weight, dim=1) + denom = torch.sum(weight * weight, dim=1) + rq = numerator / denom + if method == "alignment" and relative: + rq = rq / torch.trace(cc) + return rq + else: + # user should call patchwise_alignment or do proper flatten + raise ValueError( + f"alignment() got {input.ndim}-D data. For patchwise CNN, call patchwise_alignment(...). " + "For single-cov, ensure input is 2D via net._preprocess_inputs(...)." + ) + + +@torch.no_grad() +def expected_alignment_distribution(eigenvalues, relative=True, valid_rotation=True, with_rotation=True, bins=11, num_tests=100): + """ + From a set of eigenvalues, simulate random weight vectors + (optional orthonormal rotation) to measure an 'expected' alignment distribution. + """ + import math + import numpy as np + + N = len(eigenvalues) + if relative: + eigenvalues = eigenvalues / eigenvalues.sum() + eigenvalues = eigenvalues.view(-1, 1).expand(-1, N * num_tests) + + device = eigenvalues.device + if with_rotation: + if valid_rotation: + mixing = [] + for _ in range(num_tests): + mat = torch.normal(0, 1 / math.sqrt(N), (N, N)).to(device) + q, _ = torch.linalg.qr(mat) + mixing.append(q.T) + coefficients = torch.cat(mixing, dim=1) ** 2 + else: + coefficients = torch.normal(0, 1 / math.sqrt(N), (N, N * num_tests)).to(device) ** 2 + else: + coefficients = torch.ones((N, N * num_tests), device=device) + + weights = eigenvalues * coefficients + align = torch.sum(eigenvalues * weights, dim=0) / weights.sum(dim=0) + align_np = align.cpu().numpy() + + counts, bin_edges = np.histogram(align_np, bins=bins, density=True) + centers = (bin_edges[:-1] + bin_edges[1:]) / 2.0 + return (torch.tensor(counts, dtype=torch.float), + torch.tensor(bin_edges, dtype=torch.float), + torch.tensor(centers, dtype=torch.float)) \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/alignment_metrics_rem.py b/_arxiv/_archive/alignment_preref/alignment_metrics_rem.py new file mode 100644 index 00000000..eafbd8bb --- /dev/null +++ b/_arxiv/_archive/alignment_preref/alignment_metrics_rem.py @@ -0,0 +1,43 @@ +""" +Alignment metrics for neural network analysis. + +DEPRECATED: This module is deprecated and will be removed in a future version. +Please import directly from `alignment.metrics` instead. + +This module provides various metrics for measuring alignment between weight vectors +and activation vectors, including RQ (representation quality), MI (mutual information), +and other metrics for analyzing neural network representations. +""" + +import warnings +from typing import Dict, List, Tuple, Union, Optional, Callable + +# Show deprecation warning when the module is imported +warnings.warn( + "The alignment.alignment_metrics module is deprecated and will be removed in a future version. " + "Please import directly from alignment.metrics instead.", + DeprecationWarning, + stacklevel=2 +) + +# Import from metrics_utils instead of metrics.py to avoid circular dependencies +from alignment.utils.metrics_utils import ( + AlignmentMetricBase, + RQMetric, + MIMetric, + WeightSimilarityMetric, + NodeRedundancyMetric, + AlignmentMetricsFactory as AlignmentMetrics, + alignment +) + +# For backward compatibility, re-export everything +__all__ = [ + 'AlignmentMetricBase', + 'RQMetric', + 'MIMetric', + 'WeightSimilarityMetric', + 'NodeRedundancyMetric', + 'AlignmentMetrics', + 'alignment' +] \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/config.py b/_arxiv/_archive/alignment_preref/config.py new file mode 100644 index 00000000..b80eef6a --- /dev/null +++ b/_arxiv/_archive/alignment_preref/config.py @@ -0,0 +1,116 @@ +# -------------------------------------------- +# config.py +# -------------------------------------------- +from os import PathLike +from dataclasses import dataclass, field +from typing import (cast, List, Union, Type, TypeVar, Optional) +from omegaconf import OmegaConf as om +from omegaconf.errors import OmegaConfBaseException + +C = TypeVar("C", bound="BaseConfig") + +class BaseConfig: + @classmethod + def load(cls: Type[C], path: Union[str, PathLike]) -> C: + schema = om.structured(cls) + try: + raw = om.load(str(path)) + conf = om.merge(schema, raw) + return cast(C, om.to_object(conf)) + except OmegaConfBaseException as e: + raise ValueError(str(e)) + +@dataclass +class ModelConfig(BaseConfig): + name: str = "MLP" + dropout: float = 0.0 + alignment_layers: Optional[dict] = None + +@dataclass +class DatasetConfig(BaseConfig): + name: str = "MNIST" + path: Optional[str] = None + +@dataclass +class TrainingConfig(BaseConfig): + do_train: bool = True + epochs: int = 10 + batch_size: int = 128 + replicates: int = 1 + name: str = "Adam" + lr: float = 1e-3 + weight_decay: float = 0.0 + +@dataclass +class AlignmentConfig(BaseConfig): + do_alignment: bool = True + methods: List[str] = field(default_factory=lambda: ["RQ"]) + measure_expected: bool = True + frequency: int = 1 + bins: int = 50 + cnn_mode: str = "unfold" + scale_by_norm: bool = False # If True, scales covariance matrices by their norm before RQ calculation + +@dataclass +class CheckpointingConfig(BaseConfig): + use_prev: bool = False + save_checkpoints: bool = False + frequency: int = 1 + use_wandb: bool = False + +@dataclass +class PlotsConfig(BaseConfig): + show_loss: bool = True + show_dropout: bool = True + show_eigenfeatures: bool = True + show_eig_dropout: bool = True + +@dataclass +class ExtraConfig(BaseConfig): + num_drops: int = 9 + progressive_dropout_on_train: bool = False + single_layer_mode: bool = False + aggregate_alignment: bool = False + dropout_pruning_mode: str = "global" # "global", "per_layer_combined", or "per_layer_independent" + dropout_mode: str = "scaled" # "scaled" or "unscaled" - how targeted dropout is applied to neurons + exclude_classification_layer: bool = False # If True, excludes the classification layer from pruning + +@dataclass +class ExperimentConfig(BaseConfig): + experiment: Optional[str] = None + results_path: Optional[str] = None + no_save: bool = False + just_plot: bool = False + save_networks: bool = False + show_params: bool = False + show_all: bool = False + device: Optional[str] = None + use_timestamp: bool = False + timestamp: Optional[str] = None + ddp_world_size: int = 1 + + dataset: DatasetConfig = field(default_factory=DatasetConfig) + model: ModelConfig = field(default_factory=ModelConfig) + training: TrainingConfig = field(default_factory=TrainingConfig) + alignment: AlignmentConfig = field(default_factory=AlignmentConfig) + checkpointing: CheckpointingConfig = field(default_factory=CheckpointingConfig) + plots: PlotsConfig = field(default_factory=PlotsConfig) + extra: ExtraConfig = field(default_factory=ExtraConfig) + +@dataclass +class ExperimentArgs: + """Experiment arguments.""" + dataclass_fields = [] + + exp_name: str + name: str + device: str = "cpu" + + +@dataclass +class ExtraArgs: + """Arguments for various experiments.""" + aggregate_alignment: bool = False + num_drops: int = 9 + dropout_pruning_mode: str = "global" # "global", "per_layer_combined", or "per_layer_independent" + exclude_classification_layer: bool = False # If True, excludes the classification layer from pruning \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/datasets.py b/_arxiv/_archive/alignment_preref/datasets.py new file mode 100644 index 00000000..4aab900a --- /dev/null +++ b/_arxiv/_archive/alignment_preref/datasets.py @@ -0,0 +1,280 @@ +# -------------------------------------------- +# datasets.py +# -------------------------------------------- + +import sys +from pathlib import Path +from warnings import warn +from abc import ABC, abstractmethod + +import torch +import torchvision +from torch import nn +from torch.utils.data.distributed import DistributedSampler +from torchvision.transforms import v2 as transforms + +from alignment.models.base import AlignmentNetwork +from alignment.config import ExperimentConfig + +REQUIRED_PROPERTIES = ["dataset_constructor", "loss_function"] + +def default_loader_parameters( + distributed, + batch_size=1024, + num_workers=2, + shuffle=True, + pin_memory=True, + persistent_workers=True, +): + """ + contains the default dataloader parameters with the option of updating them + using key word arguments + + # adjusting 'shuffle' when using distributed data parallel. + """ + default_parameters = dict( + batch_size=batch_size, + num_workers=num_workers, + shuffle=False if distributed else shuffle, + pin_memory=pin_memory, + persistent_workers=persistent_workers, + ) + return default_parameters + +class DataSet(ABC): + """ + Abstract base class for any dataset wrapper. + Responsible for creating train/test datasets & loaders, + storing transforms, etc. + """ + + def __init__( + self, + device=None, + distributed=False, + dataset_parameters={}, + transform_parameters={}, + loader_parameters={}, + ): + self.set_properties() + self.check_properties() + + self.device = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") + self.distributed = distributed + + self.extra_transform = transform_parameters.pop("extra_transform", None) + + self.transform_parameters = transform_parameters + self.make_transform(**transform_parameters) + + self.dataloader_parameters = default_loader_parameters(distributed, **loader_parameters) + self.dataset_parameters = dataset_parameters + self.load_dataset(**dataset_parameters) + + def check_properties(self): + if not all([hasattr(self, prop) for prop in REQUIRED_PROPERTIES]): + not_found = [prop for prop in REQUIRED_PROPERTIES if not hasattr(self, prop)] + raise ValueError(f"The following required properties were not set: {not_found}") + + @abstractmethod + def set_properties(self): + """ + Must set: + self.dataset_constructor + self.loss_function + and optionally self.dist_params for normalization, etc. + """ + pass + + @abstractmethod + def dataset_kwargs(self, train=True, **kwargs): + """ + Returns the dict of kwargs for constructing the train or test dataset. + """ + pass + + def load_dataset(self, **kwargs): + """Instantiate self.train_dataset and self.test_dataset, then build DataLoaders.""" + self.train_dataset = self.dataset_constructor(**self.dataset_kwargs(train=True, **kwargs)) + self.test_dataset = self.dataset_constructor(**self.dataset_kwargs(train=False, **kwargs)) + self.train_sampler = DistributedSampler(self.train_dataset) if self.distributed else None + self.test_sampler = DistributedSampler(self.test_dataset) if self.distributed else None + self.train_loader = torch.utils.data.DataLoader(self.train_dataset, sampler=self.train_sampler, **self.dataloader_parameters) + self.test_loader = torch.utils.data.DataLoader(self.test_dataset, sampler=self.test_sampler, **self.dataloader_parameters) + + def unwrap_batch(self, batch, device=None): + """ + Simple method for unwrapping a batch (inputs, targets) + and placing them on the correct device. + """ + device = self.device if device is None else device + if self.extra_transform: + if isinstance(self.extra_transform, list): + for et in self.extra_transform: + batch = et(batch) + else: + warn("extra_transform is not a list, this is deprecated!", DeprecationWarning, stacklevel=2) + batch = self.extra_transform(batch) + inputs, targets = batch + inputs, targets = inputs.to(device), targets.to(device) + return inputs, targets + + def make_transform(self, center_crop=None, resize=None, flatten=False, out_channels=None): + """ + Create a transforms.Compose for the dataset. + Default: converts to float32, normalizes, etc. + + # For standard dataset usage, we typically do: + # transforms.ToImage(), + # transforms.ToDtype(torch.float32, scale=True), + # transforms.Normalize(...), + # etc. + """ + use_transforms = [ + transforms.ToImage(), + transforms.ToDtype(torch.float32, scale=True), + ] + if center_crop: + use_transforms.append(transforms.CenterCrop(center_crop)) + use_transforms.append(transforms.Normalize((self.dist_params["mean"]), (self.dist_params["std"]))) + if resize: + use_transforms.append(transforms.Resize(resize, antialias=True)) + if out_channels: + use_transforms.append(transforms.Grayscale(num_output_channels=out_channels)) + if flatten: + use_transforms.append(transforms.Lambda(torch.flatten)) + self.transform = transforms.Compose(use_transforms) + + def measure_loss(self, outputs, targets, reduction=None): + """Compute the loss via self.loss_function, optionally changing reduction.""" + if reduction is None: + return self.loss_function(outputs, targets) + standard_reduction = self.loss_function.reduction + self.loss_function.reduction = reduction + loss = self.loss_function(outputs, targets) + self.loss_function.reduction = standard_reduction + return loss + + def measure_accuracy(self, outputs, targets, k=1, percentage=True): + """ + Return top-k accuracy on classification problems. + By default, top-1 accuracy in percentage form. + """ + # Use a simple, direct implementation that's guaranteed to work correctly + # First check for invalid values + if torch.isnan(outputs).any() or torch.isinf(outputs).any(): + print("WARNING: NaN or Inf values in measure_accuracy inputs") + return torch.tensor(0.0, device=outputs.device) + + # Get the predicted classes (most likely class for each sample) + _, predicted = outputs.max(1) + + # Calculate accuracy + correct = (predicted == targets).sum().item() + total = targets.size(0) + + # Convert to percentage if requested + accuracy = (correct / total) * (100.0 if percentage else 1.0) + + + return torch.tensor(accuracy, device=outputs.device) + + +class MNIST(DataSet): + def set_properties(self): + """defines the required properties for MNIST""" + self.dataset_constructor = torchvision.datasets.MNIST + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.1307], std=[0.3081]) + + def dataset_kwargs(self, train=True, download=False, root=None): + kwargs = dict( + train=train, + download=download, + root=root, + transform=self.transform, + ) + return kwargs + +class CIFAR10(DataSet): + def set_properties(self): + self.dataset_constructor = torchvision.datasets.CIFAR10 + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + def dataset_kwargs(self, train=True, download=False, root=None): + kwargs = dict( + train=train, + download=download, + root=root, + transform=self.transform, + ) + return kwargs + +class CIFAR100(CIFAR10): + def set_properties(self): + self.dataset_constructor = torchvision.datasets.CIFAR100 + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + +class ImageNet2012(DataSet): + def set_properties(self): + """ + defines the required properties for ImageNet 2012 + (ILSVRC2012) with 1000 classes. + """ + self.dataset_constructor = torchvision.datasets.ImageNet + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + self.center_crop = 224 + + def dataset_kwargs(self, train=True, root=None): + kwargs = dict( + split="train" if train else "val", + root=root, + transform=self.transform, + ) + return kwargs + +DATASET_REGISTRY = { + "MNIST": MNIST, + "CIFAR10": CIFAR10, + "CIFAR100": CIFAR100, + "ImageNet": ImageNet2012, +} + +def get_dataset( + dataset_name, + build=False, + dataset_parameters={}, + transform_parameters={}, + loader_parameters={}, + **kwargs, +): + """ + lookup dataset constructor from dataset registry by name. + if build=True, return an instantiated dataset object, + otherwise return the constructor. + """ + if dataset_name not in DATASET_REGISTRY: + raise ValueError(f"Dataset ({dataset_name}) is not in DATASET_REGISTRY") + dataset = DATASET_REGISTRY[dataset_name] + if build: + if not isinstance(transform_parameters, dict): + raise TypeError("transform_parameters must be a dictionary") + return dataset( + dataset_parameters=dataset_parameters, + transform_parameters=transform_parameters, + loader_parameters=loader_parameters, + **kwargs, + ) + return dataset + +if __name__ == "__main__": + try: + yaml_path, args_list = sys.argv[1] + except IndexError: + raise ValueError(f"Usage: {sys.argv[0]} [CONFIG_PATH]") + + cfg = ExperimentConfig.load(yaml_path) + dataset = get_dataset(cfg.dataset.name, build=True, dataset_parameters=dict(download=True, root=Path(cfg.dataset.path))) diff --git a/_arxiv/_archive/alignment_preref/plotting.py b/_arxiv/_archive/alignment_preref/plotting.py new file mode 100644 index 00000000..ff683467 --- /dev/null +++ b/_arxiv/_archive/alignment_preref/plotting.py @@ -0,0 +1,861 @@ +# -------------------------------------------- +# plotting.py +# -------------------------------------------- + +import torch +import numpy as np +from tqdm import tqdm +from matplotlib import pyplot as plt +import matplotlib as mpl + +from alignment.utils import compute_stats_by_type, named_transpose, transpose_list, rms + + +def plot_train_results(exp, train_results, test_results, prms): + """ + Plot training curves for loss and accuracy, plus test performance. + Optionally plots alignment across training if present in `train_results`. + + If 'loss' or 'accuracy' are missing in train_results (e.g. do_train=false), + we skip the training curve plot. We can still show test performance if present. + """ + + import torch + import numpy as np + import matplotlib.pyplot as plt + import matplotlib as mpl + from alignment.utils import compute_stats_by_type + + # If do_train was false, we might not have train_results["loss"] or train_results["accuracy"]. + # We'll do a check: + has_train_loss = "loss" in train_results and train_results["loss"].dim() == 2 + has_train_acc = "accuracy" in train_results and train_results["accuracy"].dim() == 2 + + # In normal training, shape is (num_replicates, num_epochs) + # If missing, we can skip them. We'll define num_train_epochs = 0 if missing + if has_train_loss: + num_train_epochs = train_results["loss"].size(1) + else: + num_train_epochs = 0 + + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + + print("getting statistics on run data...") + plot_alignment = "alignment" in train_results + alignment = None + if plot_alignment: + alignment_by_method = {} + for record in train_results["alignment"]: + if isinstance(record, torch.Tensor): + method_key = 'RQ' + if record.dim() == 3: + net_avg = record.mean(dim=0) + elif record.dim() == 4: + net_avg = record.squeeze(0).mean(dim=0) + else: + raise ValueError(f"Unexpected tensor shape: {record.shape}") + for layer_idx in range(net_avg.size(0)): + if method_key not in alignment_by_method: + alignment_by_method[method_key] = {} + if layer_idx not in alignment_by_method[method_key]: + alignment_by_method[method_key][layer_idx] = [] + alignment_by_method[method_key][layer_idx].append(torch.mean(net_avg[layer_idx]).item()) + + elif isinstance(record, dict): + data_field = record["data"] + if isinstance(data_field, list): + for net_data in data_field: + for layer_idx, layer_list in enumerate(net_data): + # FIX: layer_list was a list, so we iterate over it before calling .items() + if isinstance(layer_list, dict): + for method_key, tensor in layer_list.items(): + if method_key not in alignment_by_method: + alignment_by_method[method_key] = {} + if layer_idx not in alignment_by_method[method_key]: + alignment_by_method[method_key][layer_idx] = [] + alignment_by_method[method_key][layer_idx].append(torch.mean(tensor).item()) + elif isinstance(layer_list, list): + for sub_dict in layer_list: + for method_key, tensor in sub_dict.items(): + if method_key not in alignment_by_method: + alignment_by_method[method_key] = {} + if layer_idx not in alignment_by_method[method_key]: + alignment_by_method[method_key][layer_idx] = [] + alignment_by_method[method_key][layer_idx].append(torch.mean(tensor).item()) + else: + raise TypeError(f"layer_idx={layer_idx} has unexpected type {type(layer_list)}") + elif isinstance(data_field, torch.Tensor): + method_key = 'RQ' + if data_field.dim() == 3: + net_avg = data_field.mean(dim=0) + elif data_field.dim() == 4: + net_avg = data_field.squeeze(0).mean(dim=0) + else: + raise ValueError(f"Unexpected record['data'] shape: {data_field.shape}") + for layer_idx in range(net_avg.size(0)): + if method_key not in alignment_by_method: + alignment_by_method[method_key] = {} + if layer_idx not in alignment_by_method[method_key]: + alignment_by_method[method_key][layer_idx] = [] + alignment_by_method[method_key][layer_idx].append(torch.mean(net_avg[layer_idx]).item()) + else: + raise TypeError(f"record['data'] has unexpected type {type(data_field)}") + else: + raise TypeError(f"train_results['alignment'] has an element of type {type(record)}") + + alignment_avg = {} + for method_key, layers_dict in alignment_by_method.items(): + layer_avgs = [] + for layer_idx in sorted(layers_dict.keys()): + vals = torch.tensor(layers_dict[layer_idx]) + layer_avgs.append(torch.mean(vals)) + alignment_avg[method_key] = torch.stack(layer_avgs, dim=0) + alignment = alignment_avg + + cmap = mpl.colormaps["tab10"] + + if has_train_loss: + train_loss_mean, train_loss_se = compute_stats_by_type( + train_results["loss"].transpose(0,1), + num_types=num_types, + dim=1, + method="se" + ) + else: + train_loss_mean, train_loss_se = None, None + + if has_train_acc: + train_acc_mean, train_acc_se = compute_stats_by_type( + train_results["accuracy"].transpose(0,1), + num_types=num_types, + dim=1, + method="se" + ) + else: + train_acc_mean, train_acc_se = None, None + + test_loss_mean, test_loss_se = None, None + test_acc_mean, test_acc_se = None, None + if ("loss" in test_results) and (test_results["loss"].dim() == 1): + test_loss_mean, test_loss_se = compute_stats_by_type( + test_results["loss"], + num_types=num_types, + dim=0, + method="se" + ) + if ("accuracy" in test_results) and (test_results["accuracy"].dim() == 1): + test_acc_mean, test_acc_se = compute_stats_by_type( + test_results["accuracy"], + num_types=num_types, + dim=0, + method="se" + ) + + print("plotting run data...") + xOffset = [-0.2, 0.2] + get_x = lambda idx: [xOffset[0] + idx, xOffset[1] + idx] + + alpha = 0.3 + figdim = 3 + figratio = 2 + width_ratios = [figdim, figdim / figratio, figdim, figdim / figratio] + fig, ax = plt.subplots(1, 4, figsize=(sum(width_ratios), figdim), width_ratios=width_ratios, layout="constrained") + + for idx, label in enumerate(labels): + if train_loss_mean is not None: + cmn = train_loss_mean[:, idx].cpu() + cse = train_loss_se[:, idx].cpu() + # Use 1-based epoch numbering for the x-axis + epochs = [i+1 for i in range(num_train_epochs)] + ax[0].plot(epochs, cmn, color=cmap(idx), label=label) + ax[0].fill_between(epochs, cmn + cse, cmn - cse, color=(cmap(idx), alpha)) + + if test_loss_mean is not None: + tmn = test_loss_mean[idx].cpu().item() + tse = test_loss_se[idx].cpu().item() + ax[1].plot(get_x(idx), [tmn]*2, color=cmap(idx), label=label, lw=4) + ax[1].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5) + + ax[0].set_xlabel("Training Epoch") + ax[0].set_ylabel("Loss") + ax[0].set_title("Training Loss") + ax[1].set_title("Testing") + + if train_loss_mean is not None: + ax[0].set_ylim(0, None) + # Set x-axis limits to show full range of epochs, with a minimum range + if num_train_epochs > 1: + ax[0].set_xlim(1, num_train_epochs) + else: + # For single epoch case, add some padding + ax[0].set_xlim(0.5, 1.5) + ylims = ax[0].get_ylim() + ax[1].set_ylim(ylims) + ax[1].set_ylabel("Loss") + ax[1].set_xlim(-0.5, num_types - 0.5) + ax[1].set_xticks(range(num_types)) + ax[1].set_xticklabels(labels, rotation=45, ha="right", fontsize=8) + + for idx, label in enumerate(labels): + if train_acc_mean is not None: + cmn = train_acc_mean[:, idx].cpu() + cse = train_acc_se[:, idx].cpu() + # Use 1-based epoch numbering for the x-axis (same as loss plot) + epochs = [i+1 for i in range(num_train_epochs)] + ax[2].plot(epochs, cmn, color=cmap(idx), label=label) + ax[2].fill_between(epochs, cmn + cse, cmn - cse, color=(cmap(idx), alpha)) + + if test_acc_mean is not None: + tmn = test_acc_mean[idx].cpu().item() + tse = test_acc_se[idx].cpu().item() + ax[3].plot(get_x(idx), [tmn]*2, color=cmap(idx), label=label, lw=4) + ax[3].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5) + + ax[2].set_xlabel("Training Epoch") + ax[2].set_ylabel("Accuracy (%)") + ax[2].set_title("Training Accuracy") + ax[3].set_title("Testing") + if train_acc_mean is not None: + ax[2].set_ylim(0, 100) + # Set x-axis limits to show full range of epochs, with a minimum range + if num_train_epochs > 1: + ax[2].set_xlim(1, num_train_epochs) + else: + # For single epoch case, add some padding + ax[2].set_xlim(0.5, 1.5) + ax[3].set_ylim(0, 100) + ax[3].set_xlim(-0.5, num_types - 0.5) + ax[3].set_xticks(range(num_types)) + ax[3].set_xticklabels(labels, rotation=45, ha="right", fontsize=8) + + exp.plot_ready("train_test_performance") + + if plot_alignment and alignment is not None: + for method_key, align_tensor in alignment.items(): + num_layers = align_tensor.size(0) + fig2, ax2 = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", sharex=True, squeeze=False) + for layer in range(num_layers): + val = align_tensor[layer].cpu().item() * 100 + ax2[0, layer].axhline(val, color=cmap(0), label=method_key) + ax2[0, layer].set_xlabel("Snapshots") + ax2[0, layer].set_ylabel("Alignment (%)") + ax2[0, layer].set_title(f"Layer {layer} - {method_key}") + exp.plot_ready(f"train_alignment_{method_key}") + + +def plot_dropout_results(exp, dropout_results, dropout_parameters, prms, dropout_type="nodes"): + """ + Plot progressive dropout results for nodes (or other types). + """ + if not exp.args.plots.show_dropout: + return + + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + cmap = mpl.colormaps["Set1"] + alpha = 0.3 + msize = 10 + figdim = 3 + + # Retrieve the arrays from `dropout_results`: + loss_high = dropout_results["progdrop_loss_high"] # shape => (N, D, ...) 3D or 4D + loss_low = dropout_results["progdrop_loss_low"] + loss_rand = dropout_results["progdrop_loss_rand"] + acc_high = dropout_results["progdrop_acc_high"] + acc_low = dropout_results["progdrop_acc_low"] + acc_rand = dropout_results["progdrop_acc_rand"] + + dropout_fraction = dropout_results["dropout_fraction"] # shape => (D,) + pruning_mode = dropout_results.get("pruning_mode", "global") # Use pruning_mode instead of by_layer + is_per_layer = pruning_mode != "global" # Check if using per-layer modes + extra_name = (pruning_mode + "_" if pruning_mode else "") + dropout_type + + # For per_layer_combined mode, use the combined results to create a single figure + if pruning_mode == "per_layer_combined" and "combined_progdrop_loss_high" in dropout_results: + print("Using combined results for per_layer_combined mode...") + loss_high = dropout_results["combined_progdrop_loss_high"] + loss_low = dropout_results["combined_progdrop_loss_low"] + loss_rand = dropout_results["combined_progdrop_loss_rand"] + acc_high = dropout_results["combined_progdrop_acc_high"] + acc_low = dropout_results["combined_progdrop_acc_low"] + acc_rand = dropout_results["combined_progdrop_acc_rand"] + is_per_layer = False # Treat as a global figure with single layer results + extra_name = "per_layer_combined_aggregated_" + dropout_type + + print("Computing dropout stats over replicate dimension...") + # shape => (N, D, E?, L?) => We do `dim=0` to average across replicates + mean_loss_high, se_loss_high = compute_stats_by_type(loss_high, num_types=num_types, dim=0, method="se") + mean_loss_low, se_loss_low = compute_stats_by_type(loss_low, num_types=num_types, dim=0, method="se") + mean_loss_rand, se_loss_rand= compute_stats_by_type(loss_rand, num_types=num_types, dim=0, method="se") + + mean_acc_high, se_acc_high = compute_stats_by_type(acc_high, num_types=num_types, dim=0, method="se") + mean_acc_low, se_acc_low = compute_stats_by_type(acc_low, num_types=num_types, dim=0, method="se") + mean_acc_rand, se_acc_rand= compute_stats_by_type(acc_rand, num_types=num_types, dim=0, method="se") + + # Now each is either 3D or 4D. We'll detect: + # e.g. mean_loss_high.shape => 3D: (num_types, D, L) + # 4D: (num_types, D, E, L) + shape_ = mean_loss_high.shape + # shape_[0] = num_types, shape_[1] = D + # we check len(shape_): + if len(shape_) == 3: + # interpret => (num_types, D, L) + # => no epoch dimension => treat E=1 + # We can reshape => (num_types, D, 1, L) to unify logic + mean_loss_high = mean_loss_high.unsqueeze(2) # => (num_types, D, 1, L) + se_loss_high = se_loss_high.unsqueeze(2) + mean_loss_low = mean_loss_low.unsqueeze(2) + se_loss_low = se_loss_low.unsqueeze(2) + mean_loss_rand = mean_loss_rand.unsqueeze(2) + se_loss_rand = se_loss_rand.unsqueeze(2) + + mean_acc_high = mean_acc_high.unsqueeze(2) + se_acc_high = se_acc_high.unsqueeze(2) + mean_acc_low = mean_acc_low.unsqueeze(2) + se_acc_low = se_acc_low.unsqueeze(2) + mean_acc_rand = mean_acc_rand.unsqueeze(2) + se_acc_rand = se_acc_rand.unsqueeze(2) + + shape_ = mean_loss_high.shape # now => (num_types, D, 1, L) + + # now shape_ => (num_types, D, E, L) + # parse + num_types_, num_drop, num_epochs, num_layers = shape_ + print(f"Shape => (num_types={num_types_}, #drops={num_drop}, #epochs={num_epochs}, #layers={num_layers})") + + # Try to get alignment layer names from the first net + layer_names = None + if hasattr(exp, "nets") and exp.nets and hasattr(exp.nets[0], "alignment_names"): + layer_names = exp.nets[0].alignment_names + + lines = ["From high", "From low", "Random"] + + # ---------------- + # Plot Loss + # ---------------- + fig_loss, ax_loss = plt.subplots( + num_epochs, num_layers, + figsize=(num_layers * figdim, num_epochs * figdim), + sharex=True, + sharey=True, + layout="constrained", + squeeze=False + ) + + ax_loss = np.reshape(ax_loss, (num_epochs, num_layers)) + print("Plotting dropout: loss...") + + for eidx in range(num_epochs): + for lidx in range(num_layers): + # Title with layer name if we have them + if layer_names and lidx < len(layer_names): + layer_str = layer_names[lidx] + else: + layer_str = f"Layer {lidx}" + # If we want "Epoch eidx" or "Snapshot eidx" + subplot_title = f"Epoch {eidx}, {layer_str}" + + # We'll do multiple lines: each line is one "type" + "From high/low/rand" + # Actually, we have 3 lines in one type or 3 lines total across all types? + # If you want each type to have 3 lines in the same subplot, do: + for ty_idx, label_str in enumerate(labels): + for i_line, line_name in enumerate(lines): + if line_name == "From high": + cmn = mean_loss_high[ty_idx, :, eidx, lidx] # shape => (#drop,) + cse = se_loss_high[ty_idx, :, eidx, lidx] + elif line_name == "From low": + cmn = mean_loss_low[ty_idx, :, eidx, lidx] + cse = se_loss_low[ty_idx, :, eidx, lidx] + else: # "From random" + cmn = mean_loss_rand[ty_idx, :, eidx, lidx] + cse = se_loss_rand[ty_idx, :, eidx, lidx] + + ax_loss[eidx, lidx].plot( + dropout_fraction, + cmn, + color=cmap(i_line), + marker=".", + markersize=msize, + label=f"{label_str} ({line_name})", + ) + ax_loss[eidx, lidx].fill_between( + dropout_fraction, + cmn - cse, + cmn + cse, + color=(cmap(i_line), alpha) + ) + + ax_loss[eidx, lidx].set_title(subplot_title) + if eidx == num_epochs - 1: + ax_loss[eidx, lidx].set_xlabel("Dropout Fraction") + ax_loss[eidx, lidx].set_xlim(0, 1) + if lidx == 0: + ax_loss[eidx, lidx].set_ylabel("Loss w/ Dropout") + ax_loss[eidx, lidx].legend(loc="best") + + exp.plot_ready("prog_dropout_" + extra_name + "_loss") + + # ---------------- + # Plot Accuracy + # ---------------- + fig_acc, ax_acc = plt.subplots( + num_epochs, num_layers, + figsize=(num_layers * figdim, num_epochs * figdim), + sharex=True, + sharey=True, + layout="constrained", + squeeze=False + ) + + ax_acc = np.reshape(ax_acc, (num_epochs, num_layers)) + print("Plotting dropout: accuracy...") + + for eidx in range(num_epochs): + for lidx in range(num_layers): + if layer_names and lidx < len(layer_names): + layer_str = layer_names[lidx] + else: + layer_str = f"Layer {lidx}" + subplot_title = f"Epoch {eidx}, {layer_str}" + + for ty_idx, label_str in enumerate(labels): + for i_line, line_name in enumerate(lines): + if line_name == "From high": + cmn = mean_acc_high[ty_idx, :, eidx, lidx] + cse = se_acc_high[ty_idx, :, eidx, lidx] + elif line_name == "From low": + cmn = mean_acc_low[ty_idx, :, eidx, lidx] + cse = se_acc_low[ty_idx, :, eidx, lidx] + else: # random + cmn = mean_acc_rand[ty_idx, :, eidx, lidx] + cse = se_acc_rand[ty_idx, :, eidx, lidx] + + ax_acc[eidx, lidx].plot( + dropout_fraction, + cmn, + color=cmap(i_line), + marker=".", + markersize=msize, + label=f"{label_str} ({line_name})", + ) + ax_acc[eidx, lidx].fill_between( + dropout_fraction, + cmn - cse, + cmn + cse, + color=(cmap(i_line), alpha) + ) + + ax_acc[eidx, lidx].set_title(subplot_title) + ax_acc[eidx, lidx].set_ylim(0, 100) + if eidx == num_epochs - 1: + ax_acc[eidx, lidx].set_xlabel("Dropout Fraction") + ax_acc[eidx, lidx].set_xlim(0, 1) + if lidx == 0: + ax_acc[eidx, lidx].set_ylabel("Accuracy w/ Dropout") + + ax_acc[eidx, lidx].legend(loc="best") + + exp.plot_ready("prog_dropout_" + extra_name + "_accuracy") + +def plot_eigenfeatures(exp, results, prms): + """ + method for plotting results related to eigen-analysis (beta, eigenvalues, class_betas). + """ + beta, eigvals, class_betas, class_names = ( + results["beta"], + results["eigvals"], + results["class_betas"], + results["class_names"], + ) + beta = [[torch.abs(b) for b in net_beta] for net_beta in beta] + class_betas = [[rms(cb, dim=2) for cb in net_class_beta] for net_class_beta in class_betas] + + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + cmap = mpl.colormaps["tab10"] + class_cmap = mpl.colormaps["viridis"].resampled(len(class_names)) + + print("measuring statistics of eigenfeature analyses...") + + beta = [torch.stack(b) for b in transpose_list(beta)] + eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)] + class_betas = [torch.stack(cb) for cb in transpose_list(class_betas)] + + beta = [b / b.sum(dim=2, keepdim=True) for b in beta] + eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals] + class_betas = [cb / cb.sum(dim=2, keepdim=True) for cb in class_betas] + + statprms = lambda method: dict(num_types=num_types, dim=0, method=method) + mean_evals, var_evals = named_transpose([compute_stats_by_type(ev, **statprms("var")) for ev in eigvals]) + sorted_beta = [torch.sort(b, descending=True, dim=2).values for b in beta] + mean_beta, se_beta = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in beta]) + mean_sorted, se_sorted = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in sorted_beta]) + mean_class_beta, se_class_beta = named_transpose([compute_stats_by_type(cb, **statprms("var")) for cb in class_betas]) + + print("plotting eigenfeature results...") + figdim = 3 + alpha = 0.3 + num_layers = len(mean_beta) + fig, ax = plt.subplots(2, num_layers, figsize=(num_layers * figdim, figdim * 2), layout="constrained", squeeze=False) + + for layer in range(num_layers): + num_input = mean_evals[layer].size(1) + num_nodes = mean_beta[layer].size(1) + for idx, label in enumerate(labels): + mn_ev = mean_evals[layer][idx] + se_ev = var_evals[layer][idx] + mn_beta_ = torch.mean(mean_beta[layer][idx], dim=0) + se_beta_ = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes) + mn_sort = torch.mean(mean_sorted[layer][idx], dim=0) + se_sort = torch.std(mean_sorted[layer][idx], dim=0) / np.sqrt(num_nodes) + + ax[0, layer].plot( + range(num_input), + mn_ev, + color=cmap(idx), + linestyle="--", + label="eigvals" if idx == 0 else None, + ) + ax[0, layer].plot(range(num_input), mn_beta_, color=cmap(idx), label=label) + ax[0, layer].fill_between(range(num_input), mn_beta_ + se_beta_, mn_beta_ - se_beta_, color=(cmap(idx), alpha)) + + ax[1, layer].plot(range(num_input), mn_sort, color=cmap(idx), label=label) + ax[1, layer].fill_between(range(num_input), mn_sort + se_sort, mn_sort - se_sort, color=(cmap(idx), alpha)) + + ax[0, layer].set_xscale("log") + ax[1, layer].set_xscale("log") + ax[0, layer].set_xlabel("Input Dimension") + ax[1, layer].set_xlabel("Sorted Input Dim") + ax[0, layer].set_ylabel("Relative Eigval & Beta") + ax[1, layer].set_ylabel("Relative Beta (Sorted)") + ax[0, layer].set_title(f"Layer {layer}") + ax[1, layer].set_title(f"Layer {layer}") + + if layer == num_layers - 1: + ax[0, layer].legend(loc="best") + ax[1, layer].legend(loc="best") + + exp.plot_ready("eigenfeatures") + + fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", squeeze=False) + for layer in range(num_layers): + num_input = mean_evals[layer].size(1) + num_nodes = mean_beta[layer].size(1) + for idx, label in enumerate(labels): + mn_ev = mean_evals[layer][idx] + se_ev = var_evals[layer][idx] + mn_beta_ = torch.mean(mean_beta[layer][idx], dim=0) + se_beta_ = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes) + ax[0, layer].plot( + range(num_input), + mn_ev, + color=cmap(idx), + linestyle="--", + label="eigvals" if idx == 0 else None, + ) + ax[0, layer].plot(range(num_input), mn_beta_, color=cmap(idx), label=label) + ax[0, layer].fill_between(range(num_input), mn_beta_ + se_beta_, mn_beta_ - se_beta_, color=(cmap(idx), alpha)) + + ax[0, layer].set_xscale("log") + ax[0, layer].set_yscale("log") + ax[0, layer].set_xlabel("Input Dimension") + ax[0, layer].set_ylabel("Relative Eigval & Beta") + ax[0, layer].set_title(f"Layer {layer}") + + if layer == num_layers - 1: + ax[0, layer].legend(loc="best") + + exp.plot_ready("eigenfeatures_loglog") + + fig, ax = plt.subplots( + num_types, + num_layers, + figsize=(num_layers * figdim, figdim * num_types), + layout="constrained", + squeeze=False, + ) + ax = np.reshape(ax, (num_types, num_layers)) + for layer in range(num_layers): + num_input = mean_evals[layer].size(1) + for idx, label in enumerate(labels): + for idx_class, class_name in enumerate(class_names): + mn_data = mean_class_beta[layer][idx][idx_class] + se_data = se_class_beta[layer][idx][idx_class] + ax[idx, layer].plot(range(num_input), mn_data, color=class_cmap(idx_class), label=class_name) + ax[idx, layer].fill_between( + range(num_input), + mn_data + se_data, + mn_data - se_data, + color=(class_cmap(idx_class), alpha), + ) + + ax[idx, layer].set_xscale("log") + ax[idx, layer].set_yscale("linear") + ax[idx, layer].set_xlabel("Input Dimension") + if layer == 0: + ax[idx, layer].set_ylabel(f"{label}\nClass Loading (RMS)") + if idx == 0: + ax[idx, layer].set_title(f"Layer {layer}") + if layer == num_layers - 1: + ax[idx, layer].legend(loc="upper right", fontsize=6) + + exp.plot_ready("class_eigenfeatures") + + +def plot_adversarial_results(exp, eigen_results, adversarial_results, prms): + """ + Visualize adversarial attack results (accuracy vs epsilon) and structure + of adversarial perturbations in the eigenvector basis. + """ + accuracy, beta, eigvals = ( + adversarial_results["accuracy"], + adversarial_results["betas"], + eigen_results["eigvals"], + ) + epsilons, use_sign = adversarial_results["use_sign"], adversarial_results["use_sign"] + + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + cmap = mpl.colormaps["tab10"] + + print("measuring statistics of adversarial analyses...") + + accuracy = torch.stack([torch.stack(acc) for acc in transpose_list(accuracy)]) # shape (num_epsilon, num_nets) + eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)] + beta = [torch.stack(b) for b in beta] + + beta = [b / b.sum(dim=2, keepdim=True) for b in beta] + eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals] + + statprms = lambda dim, method: dict(num_types=num_types, dim=dim, method=method) + mean_acc, se_acc = compute_stats_by_type(accuracy, **statprms(1, "var")) + + from alignment.core.utils import named_transpose + mean_beta, se_beta = named_transpose([compute_stats_by_type(b, **statprms(1, "var")) for b in beta]) + mean_evals, var_evals = named_transpose([compute_stats_by_type(ev, **statprms(0, "var")) for ev in eigvals]) + + print("plotting adversarial success results...") + figdim = 3 + alpha = 0.3 + num_layers = len(mean_beta) + fig, ax = plt.subplots(1, 1, figsize=(figdim, figdim), layout="constrained") + for idx, label in enumerate(labels): + ax.plot(epsilons, mean_acc[:, idx], color=cmap(idx), label=label) + ax.set_xlabel("Epsilon") + ax.set_ylabel("Accuracy") + ax.set_title("Adversarial Attack Success") + ax.legend(loc="best") + + exp.plot_ready("adversarial_success") + + print("plotting adversarial structure...") + fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", squeeze=False) + for layer in range(num_layers): + num_input = mean_beta[layer].size(2) + for idx, label in enumerate(labels): + ax[0, layer].plot( + range(num_input), + torch.nanmean(mean_beta[layer][:, idx], dim=0).detach(), + color=cmap(idx), + label=label, + ) + ax[0, layer].set_xscale("log") + ax[0, layer].set_xlabel("Input Dimension") + ax[0, layer].set_ylabel("Average Component of Pertubation") + ax[0, layer].set_title(f"Layer {layer}") + if layer == num_layers - 1: + ax[0, layer].legend(loc="best", fontsize=6) + + exp.plot_ready("adversarial_structure") + + +def plot_rf(rf, width, alignment=None, alignBounds=None, showRFs=None, figSize=5): + """ + Helper for visualizing receptive fields in a grid. + Optionally color-coded by alignment. + """ + if showRFs is not None: + rf = rf.reshape(rf.shape[0], -1) + idxRandom = np.random.choice(range(rf.shape[0]), showRFs, replace=False) + rf = rf[idxRandom, :] + else: + showRFs = rf.shape[0] + rf = rf.T / np.abs(rf).max(axis=1) + rf = rf.T + rf = rf.reshape(showRFs, width, width) + if alignment is not None: + cmap = mpl.cm.get_cmap("rainbow", rf.shape[0]) + cmapPeak = lambda x: cmap(x) + if alignBounds is not None: + alignment = alignment - alignBounds[0] + alignment = alignment / (alignBounds[1] - alignBounds[0]) + else: + alignment = alignment - alignment.min() + alignment = alignment / alignment.max() + + n = int(np.ceil(np.sqrt(rf.shape[0]))) + fig, axes = plt.subplots(nrows=n, ncols=n, sharex=True, sharey=True, squeeze=False) + fig.set_size_inches(figSize, figSize) + + N = 1000 + for i in tqdm(range(rf.shape[0])): + ax = axes[i // n][i % n] + if alignment is not None: + vals = np.ones((N, 4)) + cAlignment = alignment[i].numpy() + cPeak = cmapPeak(alignment[i].numpy()) + vals[:, 0] = np.linspace(0, cPeak[0], N) + vals[:, 1] = np.linspace(0, cPeak[1], N) + vals[:, 2] = np.linspace(0, cPeak[2], N) + usecmap = mpl.colors.ListedColormap(vals) + ax.imshow(rf[i], cmap=usecmap, vmin=-1, vmax=1) + else: + ax.imshow(rf[i], cmap="gray", vmin=-1, vmax=1) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_aspect("equal") + for j in range(rf.shape[0], n * n): + ax = axes[j // n][j % n] + ax.imshow(np.ones_like(rf[0]) * -1, cmap="gray", vmin=-1, vmax=1) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_aspect("equal") + fig.subplots_adjust(wspace=0.0, hspace=0.0) + return fig + + + + +def plot_grad_alignment_correlation(results, return_fig=False): + """ + If return_fig=True, returns a matplotlib Figure instead of displaying. + Expects results["grad_alignment_corr"] to be a list of dicts: + { + "epoch": int, + "batch": int, + "net": int, + "corr": { layer_name -> float or None } + } + Plots correlation vs. epoch for each layer. + """ + + if "grad_alignment_corr" not in results: + print("No grad_alignment_corr found in results") + if return_fig: + fig = plt.figure() + return fig + return + + layer_corr_data = {} + for record in results["grad_alignment_corr"]: + epoch = record["epoch"] + if not isinstance(epoch, int): + continue + for layer_name, cval in record["corr"].items(): + if cval is None: + continue + if layer_name not in layer_corr_data: + layer_corr_data[layer_name] = [] + layer_corr_data[layer_name].append((epoch, cval)) + + if not layer_corr_data: + print("No valid correlation data to plot") + if return_fig: + fig = plt.figure() + return fig + return + + fig, axes = plt.subplots(len(layer_corr_data), 1, figsize=(6, 3 * len(layer_corr_data)), sharex=True) + if not isinstance(axes, np.ndarray): + axes = [axes] + sorted_layers = sorted(layer_corr_data.keys()) + + for ax, layer_name in zip(axes, sorted_layers): + data = layer_corr_data[layer_name] + data.sort(key=lambda x: x[0]) + epochs = [d[0] for d in data] + corrs = [d[1] for d in data] + ax.plot(epochs, corrs, marker="o", label=layer_name) + ax.set_ylabel("Grad vs Align Corr") + ax.set_title(f"Layer {layer_name}") + ax.legend(loc="best") + + axes[-1].set_xlabel("Epoch") + fig.tight_layout() + + if return_fig: + return fig + else: + plt.show() + +def plot_alignment_change_correlation(results, return_fig=False): + """ + If return_fig=True, returns a matplotlib Figure instead of displaying. + Expects results["alignment"] to hold epoch-wise alignment data so we can + compute correlation of alignment(e) vs alignment(e+1). This is purely illustrative. + """ + if "alignment" not in results: + print("No alignment key in results") + if return_fig: + fig = plt.figure() + return fig + return + + # Example: collect epoch->some combined alignment vector + epoch_alignments = {} + for record in results["alignment"]: + ep = record.get("epoch", None) + if not isinstance(ep, int): + continue + if not record.get("data"): + continue + # Suppose record["data"] is list-of-nets => each net is list-of-layers => dict {method->(out_features,)} + # We'll only use the first net as an example + net_data = record["data"][0] + layer_vecs = [] + for layer_dict in net_data: + # if "RQ" in layer_dict => shape (out_features,) + if "RQ" in layer_dict: + layer_vecs.append(layer_dict["RQ"].detach().cpu()) + if layer_vecs: + combined = torch.cat(layer_vecs, dim=0) + epoch_alignments[ep] = combined + + sorted_eps = sorted(epoch_alignments.keys()) + if len(sorted_eps) < 2: + print("Not enough epochs with alignment to measure changes") + if return_fig: + fig = plt.figure() + return fig + return + + corrs = [] + xvals = [] + for i in range(len(sorted_eps) - 1): + e1 = sorted_eps[i] + e2 = sorted_eps[i+1] + a1 = epoch_alignments[e1] + a2 = epoch_alignments[e2] + if a1.shape == a2.shape: + stack = torch.stack([a1, a2], dim=0) + cc = torch.corrcoef(stack) + cval = cc[0, 1].item() + corrs.append(cval) + xvals.append((e1 + e2) / 2.0) + else: + corrs.append(float("nan")) + xvals.append((e1 + e2) / 2.0) + + fig, ax = plt.subplots(figsize=(6,4)) + ax.plot(xvals, corrs, marker="o", label="Align(e) vs Align(e+1) Corr") + ax.set_xlabel("Epoch midpoint") + ax.set_ylabel("Correlation") + ax.set_title("Alignment Change Correlation") + ax.legend(loc="best") + fig.tight_layout() + + if return_fig: + return fig + else: + plt.show() \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/plotting_rem.py b/_arxiv/_archive/alignment_preref/plotting_rem.py new file mode 100644 index 00000000..362bb61c --- /dev/null +++ b/_arxiv/_archive/alignment_preref/plotting_rem.py @@ -0,0 +1,576 @@ +""" +Plotting utilities for visualization of experiment results. + +DEPRECATED: This module is kept for backward compatibility. +Please use alignment.utils.plotting instead. + +This module provides functions for visualizing experimental results from +alignment studies, including progressive dropout analysis with different +pruning strategies and other alignment metrics. +""" + +import warnings +import logging +import math +import os +from typing import Dict, List, Optional, Tuple, Union, Any +import json + +import numpy as np +import torch +import matplotlib as mpl +import matplotlib.pyplot as plt +import seaborn as sns +from matplotlib.figure import Figure +from matplotlib.axes import Axes + +# Import the new modules for re-export +from alignment.utils.plotting import ( + plot_dropout_results, + plot_experiment_summary, + plot_dropout_comparison, + log_plots_to_wandb +) + +# Show deprecation warning on import +warnings.warn( + "The alignment.plotting module is deprecated and will be removed in a future version. " + "Please use alignment.utils.plotting instead.", + DeprecationWarning, + stacklevel=2 +) + +logger = logging.getLogger(__name__) + +# Set default figure size +plt.rcParams["figure.figsize"] = (10, 6) +# Use SVG backend for better quality +plt.rcParams["figure.dpi"] = 120 + + +def plot_pruning_experiments( + results: Dict[str, Any], + figure_path: Optional[str] = None, + title_prefix: str = "Progressive Dropout", +) -> List[str]: + """ + Plot results from pruning experiments. + + DEPRECATED: Use plot_dropout_results in alignment.utils.plotting instead. + + Args: + results: Dictionary of results + figure_path: Path to save figures + title_prefix: Prefix for plot titles + + Returns: + List of saved figure paths + """ + warnings.warn( + "plot_pruning_experiments is deprecated. Use plot_dropout_results instead.", + DeprecationWarning, + stacklevel=2 + ) + + # Extract pruning and dropout modes + pruning_mode = results.get("pruning_mode", "global_joint") + dropout_mode = results.get("dropout_mode", "global") + + # Map old modes to new ones + if pruning_mode == "global": + pruning_mode = "global_joint" + elif pruning_mode == "per_layer_combined": + pruning_mode = "layer_wise" + elif pruning_mode == "per_layer_independent": + pruning_mode = "layer_isolated" + + # Call the new function + return plot_dropout_results( + results, + figure_path=figure_path, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + title_prefix=title_prefix + ) + + +def plot_per_layer_independent( + results: Dict[str, Any], + figure_path: Optional[str] = None, + title_prefix: str = "Per-Layer Independent Pruning", +) -> List[str]: + """ + Plot results from per-layer independent pruning. + + DEPRECATED: Use plot_dropout_results with pruning_mode="layer_isolated" instead. + + Args: + results: Dictionary of results + figure_path: Path to save figures + title_prefix: Prefix for plot titles + + Returns: + List of saved figure paths + """ + warnings.warn( + "plot_per_layer_independent is deprecated. Use plot_dropout_results with pruning_mode='layer_isolated' instead.", + DeprecationWarning, + stacklevel=2 + ) + + # Force pruning mode to be layer_isolated (formerly per_layer_independent) + results["pruning_mode"] = "layer_isolated" + + # Call the main plotting function + return plot_dropout_results( + results, + figure_path=figure_path, + pruning_mode="layer_isolated", + dropout_mode=results.get("dropout_mode", "global"), + title_prefix=title_prefix + ) + + +def plot_dropout_results( + results: Dict, + plot_dir: str, + pruning_mode: str = "global", + dropout_mode: str = "scaled", + title_prefix: str = "Dropout" +) -> Dict[str, str]: + """ + Plot dropout experiment results using a style similar to the preref implementation. + + Args: + results: Results dictionary from progressive_dropout + plot_dir: Directory to save plots to + pruning_mode: Pruning mode used in the experiment + dropout_mode: Dropout mode used in the experiment + title_prefix: Prefix for plot titles + + Returns: + Dictionary mapping plot types to file paths + """ + # Create the output directory if it doesn't exist + os.makedirs(plot_dir, exist_ok=True) + + # Check for error in results + if "error" in results: + logger.error(f"Cannot plot results due to error: {results['error']}") + return {} + + # Extract dropout fractions + dropout_fractions = results.get("dropout_fractions", []) + if not isinstance(dropout_fractions, (list, np.ndarray)) or len(dropout_fractions) == 0: + logger.error("No dropout fractions found in results") + return {} + + # Set up figure parameters + strategies = ["high_rq", "low_rq", "random"] + colors = {"high_rq": "blue", "low_rq": "red", "random": "green"} + linestyles = {"high_rq": "-", "low_rq": "-", "random": "-"} + markers = {"high_rq": "o", "low_rq": "s", "random": "^"} + strategy_labels = {"high_rq": "High RQ", "low_rq": "Low RQ", "random": "Random"} + + # Create dictionary to store file paths + file_paths = {} + + # ---------------- + # Plot Accuracy (single figure with all strategies) + # ---------------- + fig, ax = plt.subplots(figsize=(10, 6)) + + for strategy in strategies: + accs = results.get("accuracies", {}).get(strategy) + if not isinstance(accs, (list, np.ndarray)) or len(accs) == 0: + logger.warning(f"No accuracy data for strategy {strategy}") + continue + + ax.plot( + dropout_fractions, + accs, + marker=markers[strategy], + linestyle=linestyles[strategy], + color=colors[strategy], + linewidth=2, + markersize=6, + label=strategy_labels[strategy] + ) + + ax.set_xlabel('Dropout Fraction', fontsize=14) + ax.set_ylabel('Accuracy (%)', fontsize=14) + ax.set_title(f'{title_prefix} - {pruning_mode.replace("_", " ").title()} Pruning ({dropout_mode})', fontsize=16) + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=12) + ax.set_ylim([0, 100]) # Fixed y-axis for accuracy + + accuracy_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_accuracy.png") + plt.savefig(accuracy_file, dpi=300, bbox_inches='tight') + plt.close(fig) + + file_paths["accuracy"] = accuracy_file + logger.info(f"Saved accuracy plot to {accuracy_file}") + + # ---------------- + # Plot Loss (single figure with all strategies) + # ---------------- + fig, ax = plt.subplots(figsize=(10, 6)) + + for strategy in strategies: + losses = results.get("losses", {}).get(strategy) + if not isinstance(losses, (list, np.ndarray)) or len(losses) == 0: + logger.warning(f"No loss data for strategy {strategy}") + continue + + ax.plot( + dropout_fractions, + losses, + marker=markers[strategy], + linestyle=linestyles[strategy], + color=colors[strategy], + linewidth=2, + markersize=6, + label=strategy_labels[strategy] + ) + + ax.set_xlabel('Dropout Fraction', fontsize=14) + ax.set_ylabel('Loss (%)', fontsize=14) + ax.set_title(f'{title_prefix} Loss - {pruning_mode.replace("_", " ").title()} Pruning ({dropout_mode})', fontsize=16) + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=12) + + loss_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_loss.png") + plt.savefig(loss_file, dpi=300, bbox_inches='tight') + plt.close(fig) + + file_paths["loss"] = loss_file + logger.info(f"Saved loss plot to {loss_file}") + + # ---------------- + # Plot Individual Strategy Comparisons + # ---------------- + for strategy in strategies: + accs = results.get("accuracies", {}).get(strategy) + if not isinstance(accs, (list, np.ndarray)) or len(accs) == 0: + continue + + fig, ax = plt.subplots(figsize=(10, 6)) + + ax.plot( + dropout_fractions, + accs, + marker=markers[strategy], + linestyle=linestyles[strategy], + color=colors[strategy], + linewidth=2, + markersize=6, + label=f"{strategy_labels[strategy]} Accuracy" + ) + + # Add loss to the same plot if available + losses = results.get("losses", {}).get(strategy) + if isinstance(losses, (list, np.ndarray)) and len(losses) > 0: + ax.plot( + dropout_fractions, + losses, + marker=markers[strategy], + linestyle='--', + color=colors[strategy], + linewidth=2, + markersize=6, + label=f"{strategy_labels[strategy]} Loss" + ) + + ax.set_xlabel('Dropout Fraction', fontsize=14) + ax.set_ylabel('Percentage (%)', fontsize=14) + ax.set_title(f'{title_prefix} - {strategy_labels[strategy]} ({dropout_mode})', fontsize=16) + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=12) + + strategy_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_{strategy}.png") + plt.savefig(strategy_file, dpi=300, bbox_inches='tight') + plt.close(fig) + + file_paths[f"{strategy}_comparison"] = strategy_file + logger.info(f"Saved {strategy} comparison plot to {strategy_file}") + + # ---------------- + # Plot Alignment Values (if available) + # ---------------- + alignment_values = results.get("alignment_values", {}) + if alignment_values: + # Create one plot per layer + for layer_idx in range(len(alignment_values.get("high_rq", [{}])[0] or [])): + fig, ax = plt.subplots(figsize=(10, 6)) + + for strategy in strategies: + # Extract alignment values for this layer across dropout fractions + if strategy in alignment_values: + layer_alignments = [] + for frac_idx, alignment in enumerate(alignment_values[strategy]): + if alignment and layer_idx < len(alignment): + layer_alignments.append(alignment[layer_idx]) + else: + layer_alignments.append(None) + + # Filter out None values + valid_fractions = [] + valid_alignments = [] + for frac_idx, alignment in enumerate(layer_alignments): + if alignment is not None: + valid_fractions.append(dropout_fractions[frac_idx]) + valid_alignments.append(alignment) + + if valid_fractions and valid_alignments: + ax.plot( + valid_fractions, + valid_alignments, + marker=markers[strategy], + linestyle=linestyles[strategy], + color=colors[strategy], + linewidth=2, + markersize=6, + label=strategy_labels[strategy] + ) + + ax.set_xlabel('Dropout Fraction', fontsize=14) + ax.set_ylabel('Alignment Value', fontsize=14) + ax.set_title(f'Layer {layer_idx+1} Alignment - {pruning_mode.replace("_", " ").title()}', fontsize=16) + ax.grid(True, linestyle='--', alpha=0.7) + ax.legend(fontsize=12) + + alignment_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_layer{layer_idx+1}_alignment.png") + plt.savefig(alignment_file, dpi=300, bbox_inches='tight') + plt.close(fig) + + file_paths[f"layer{layer_idx+1}_alignment"] = alignment_file + logger.info(f"Saved layer {layer_idx+1} alignment plot to {alignment_file}") + + # Save raw results as JSON for further analysis + try: + # Helper function to safely convert tensors and numpy arrays to lists + def safe_convert(obj): + if isinstance(obj, np.ndarray): + return obj.tolist() + elif isinstance(obj, torch.Tensor): + return obj.cpu().tolist() + elif isinstance(obj, dict): + return {k: safe_convert(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [safe_convert(item) for item in obj] + else: + return obj + + # Create a simplified version of results for JSON serialization + json_results = { + "dropout_fractions": safe_convert(dropout_fractions), + "accuracies": safe_convert(results.get("accuracies", {})), + "losses": safe_convert(results.get("losses", {})), + } + + # Exclude alignment_values as they can be complex and not easily serializable + + json_path = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_results.json") + with open(json_path, 'w') as f: + json.dump(json_results, f, indent=2) + + file_paths["json_results"] = json_path + logger.info(f"Saved raw results to {json_path}") + except Exception as e: + logger.error(f"Error saving JSON results: {str(e)}") + + return file_paths + + +def plot_experiment_summary( + results: Dict, + plot_dir: str +) -> Dict[str, str]: + """ + Generate a comprehensive summary of experiment results. + + Args: + results: Results dictionary from the experiment + plot_dir: Directory to save plots to + + Returns: + Dictionary mapping plot types to file paths + """ + # Create the output directory if it doesn't exist + os.makedirs(plot_dir, exist_ok=True) + + # Prepare figure + fig = plt.figure(figsize=(15, 12)) + gs = fig.add_gridspec(2, 2, width_ratios=[1, 1], height_ratios=[1, 1]) + + # Panel 1: Configuration summary + ax1 = fig.add_subplot(gs[0, 0]) + ax1.axis('off') + + # Extract config info + config = results.get("config", {}) + config_text = [] + + if hasattr(config, "model") and hasattr(config.model, "model_name"): + config_text.append(f"Model: {config.model.model_name}") + + if hasattr(config, "dataset") and hasattr(config.dataset, "dataset_name"): + config_text.append(f"Dataset: {config.dataset.dataset_name}") + + if hasattr(config, "alignment") and hasattr(config.alignment, "metric"): + config_text.append(f"Alignment Metric: {config.alignment.metric}") + + if hasattr(config, "alignment"): + if hasattr(config.alignment, "dropout_min") and hasattr(config.alignment, "dropout_max"): + config_text.append(f"Dropout Range: {config.alignment.dropout_min} to {config.alignment.dropout_max}") + if hasattr(config.alignment, "dropout_steps"): + config_text.append(f"Dropout Steps: {config.alignment.dropout_steps}") + + if hasattr(config, "extra"): + if hasattr(config.extra, "dropout_mode"): + config_text.append(f"Dropout Mode: {config.extra.dropout_mode}") + if hasattr(config.extra, "dropout_pruning_mode"): + config_text.append(f"Dropout Pruning Mode: {config.extra.dropout_pruning_mode}") + + ax1.text(0.05, 0.95, "\n".join(config_text), fontsize=11, + verticalalignment='top', horizontalalignment='left') + ax1.set_title("Experiment Configuration", fontsize=14) + + # Panel 2: Progressive Dropout results if available + ax2 = fig.add_subplot(gs[0, 1]) + prog_results = results.get("progressive_dropout", {}) + + if "accuracies" in prog_results and "dropout_fractions" in prog_results: + fractions = prog_results["dropout_fractions"] + for strategy, color in [("high_rq", "blue"), ("low_rq", "red"), ("random", "green")]: + if strategy in prog_results["accuracies"]: + accs = prog_results["accuracies"][strategy] + ax2.plot( + fractions, + accs, + marker='o', + linestyle='-', + color=color, + label=strategy.replace('_', ' ').title() + ) + + ax2.set_xlabel("Dropout Fraction", fontsize=12) + ax2.set_ylabel("Accuracy (%)", fontsize=12) + ax2.set_title("Progressive Dropout Results", fontsize=14) + ax2.grid(True, alpha=0.3) + ax2.set_ylim([0, 105]) + ax2.legend() + else: + ax2.text(0.5, 0.5, "No Progressive Dropout Results", + fontsize=12, ha='center', va='center') + ax2.axis('off') + + # Panel 3: Eigenvector Dropout results if available + ax3 = fig.add_subplot(gs[1, 0]) + eig_results = results.get("eigenvector_dropout", {}) + + if "accuracies" in eig_results and "dropout_fractions" in eig_results: + fractions = eig_results["dropout_fractions"] + if "eigenvector" in eig_results["accuracies"]: + accs = eig_results["accuracies"]["eigenvector"] + ax3.plot( + fractions, + accs, + marker='o', + linestyle='-', + color='purple', + label="Eigenvector" + ) + + # Also add the high_rq from progressive dropout for comparison if available + if "accuracies" in prog_results and "high_rq" in prog_results["accuracies"]: + prog_fracs = prog_results["dropout_fractions"] + prog_accs = prog_results["accuracies"]["high_rq"] + # Only plot if the fractions match + if len(prog_fracs) == len(fractions) and all(a == b for a, b in zip(prog_fracs, fractions)): + ax3.plot( + fractions, + prog_accs, + marker='s', + linestyle='--', + color='blue', + label="High RQ" + ) + + ax3.set_xlabel("Dropout Fraction", fontsize=12) + ax3.set_ylabel("Accuracy (%)", fontsize=12) + ax3.set_title("Eigenvector Dropout Results", fontsize=14) + ax3.grid(True, alpha=0.3) + ax3.set_ylim([0, 105]) + ax3.legend() + else: + ax3.text(0.5, 0.5, "No Eigenvector Dropout Results", + fontsize=12, ha='center', va='center') + ax3.axis('off') + + # Panel 4: Alignment comparison or other metrics + ax4 = fig.add_subplot(gs[1, 1]) + + # Check both progressive and eigenvector results for alignment values + alignment_data = None + if "alignment_values" in prog_results and "high_rq" in prog_results["alignment_values"]: + alignment_data = prog_results["alignment_values"]["high_rq"][0] + elif "alignment_values" in eig_results and "eigenvector" in eig_results["alignment_values"]: + alignment_data = eig_results["alignment_values"]["eigenvector"][0] + + if alignment_data and len(alignment_data) > 0: + # Extract alignment values as floats + alignment_values = [] + for val in alignment_data: + if isinstance(val, torch.Tensor): + alignment_values.append(val.item()) + else: + alignment_values.append(float(val)) + + # Create bar chart of alignment by layer + x = np.arange(len(alignment_values)) + bars = ax4.bar(x, alignment_values, width=0.6, alpha=0.7) + + # Add value labels on top of bars + for i, bar in enumerate(bars): + height = bar.get_height() + ax4.text( + bar.get_x() + bar.get_width()/2., + height + 0.01, + f'{alignment_values[i]:.3f}', + ha='center', + va='bottom', + rotation=0, + fontsize=9 + ) + + ax4.set_xlabel("Layer", fontsize=12) + ax4.set_ylabel("Alignment Value", fontsize=12) + ax4.set_title("Alignment by Layer", fontsize=14) + ax4.set_xticks(x) + ax4.set_xticklabels([f"Layer {i+1}" for i in range(len(alignment_values))]) + ax4.grid(True, alpha=0.3, axis='y') + else: + ax4.text(0.5, 0.5, "No Alignment Data", + fontsize=12, ha='center', va='center') + ax4.axis('off') + + # Add an overall title + title = "Experiment Summary" + if hasattr(config, "model") and hasattr(config.model, "model_name"): + if hasattr(config, "dataset") and hasattr(config.dataset, "dataset_name"): + title += f": {config.model.model_name} on {config.dataset.dataset_name}" + else: + title += f": {config.model.model_name}" + + fig.suptitle(title, fontsize=16, y=0.98) + plt.tight_layout(rect=[0, 0, 1, 0.96]) # Adjust for suptitle + + # Save the figure + summary_file = os.path.join(plot_dir, "experiment_summary.png") + plt.savefig(summary_file, dpi=300, bbox_inches='tight') + plt.close(fig) + + logger.info(f"Saved experiment summary to {summary_file}") + + return {"summary": summary_file} \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/processing.py b/_arxiv/_archive/alignment_preref/processing.py new file mode 100644 index 00000000..a7186a89 --- /dev/null +++ b/_arxiv/_archive/alignment_preref/processing.py @@ -0,0 +1,869 @@ +# processing.py + +import os +import torch +from tqdm import tqdm + +from alignment.utils import load_checkpoints, test_nets +from Code.alignment.src.alignment_preref.alignment_metrics_rem import AlignmentMetrics +from alignment.train import train, test + +def parse_alignment_to_tensor(alignment_list, aggregate=True, by_layer=False): + """ + Convert a list of alignment records into a structure suitable for dropout sorting. + If aggregator=True => you may have multiple records per epoch or per batch. + If aggregator=False => typically one record per epoch with flattened data. + + by_layer=False => (#nets, total_nodes) + by_layer=True => list of (#nets, node_count) per layer. + """ + + if len(alignment_list) == 0: + raise ValueError("parse_alignment_to_tensor: empty alignment_list") + + if not by_layer: + all_records_tensors = [] + for record in alignment_list: + if "data" not in record: + continue + netwise_tensors = [] + # record["data"] => (#nets, #layers) + # or if aggregator=false => might be (#nets, #layers) but each layer is partial avg + for net_i_data in record["data"]: + node_tensors = [] + for layer_dict in net_i_data: + if "RQ" not in layer_dict: + raise ValueError("Expected 'RQ' in layer_dict") + node_tensors.append(layer_dict["RQ"].flatten()) + cat_tsr = torch.cat(node_tensors, dim=0) + netwise_tensors.append(cat_tsr) + if len(netwise_tensors) == 0: + continue + netwise_tsr = torch.stack(netwise_tensors, dim=0) # (#nets, total_nodes) + all_records_tensors.append(netwise_tsr) + if len(all_records_tensors) == 0: + raise ValueError("No valid alignment data found") + + bigstack = torch.stack(all_records_tensors, dim=0) # (#records, #nets, total_nodes) + if aggregate: + return bigstack.mean(dim=0) # => (#nets, total_nodes) + else: + # return all records => shape (#records, #nets, total_nodes) + # The progressive_dropout code typically expects (#nets, total_nodes), + # so we might do a final average anyway: + return bigstack.mean(dim=0) + + else: + max_layers_found = 0 + for record in alignment_list: + if "data" not in record: + continue + for net_i_data in record["data"]: + if len(net_i_data) > max_layers_found: + max_layers_found = len(net_i_data) + + layer_storage = [[] for _ in range(max_layers_found)] + + for record in alignment_list: + if "data" not in record: + continue + net_list = record["data"] # (#nets, #layers) + for layer_i in range(max_layers_found): + layer_nodevals = [] + for net_i_data in net_list: + if layer_i < len(net_i_data): + lay_dict = net_i_data[layer_i] + if "RQ" not in lay_dict: + raise ValueError("Expected 'RQ' in lay_dict") + layer_nodevals.append(lay_dict["RQ"].flatten()) + else: + layer_nodevals.append(None) + + valid_vals = [v for v in layer_nodevals if v is not None] + if len(valid_vals) == 0: + continue + node_count = valid_vals[0].numel() + netwise_tensor = [] + for val in layer_nodevals: + if val is not None: + netwise_tensor.append(val) + netwise_tensor = torch.stack(netwise_tensor, dim=0) + layer_storage[layer_i].append(netwise_tensor) + + final_layer_list = [] + for layer_i in range(max_layers_found): + if len(layer_storage[layer_i]) == 0: + final_layer_list.append(None) + continue + stacked = torch.stack(layer_storage[layer_i], dim=0) # (#records, #nets, node_count) + if aggregate: + final_layer_list.append(stacked.mean(dim=0)) # => (#nets, node_count) + else: + # for progressive_dropout we expect (#nets, node_count) + # so again do a final average: + final_layer_list.append(stacked.mean(dim=0)) + + return final_layer_list + + +def train_networks(exp, nets, optimizers, dataset, **special_parameters): + do_alignment_train = exp.args.alignment.do_alignment + methods = exp.args.alignment.methods + measure_freq = exp.args.alignment.frequency + + params = dict( + train_set=True, + num_epochs=exp.args.training.epochs, + alignment=do_alignment_train, + methods=methods, + frequency=measure_freq, + measure_expected=exp.args.alignment.measure_expected, + results=None, + verbose=True, + ) + params.update(**special_parameters) + + if exp.args.checkpointing.use_prev and os.path.isfile(exp.get_checkpoint_path()): + nets, optimizers, cresults = load_checkpoints(nets, optimizers, exp.args.device, exp.get_checkpoint_path()) + for net in nets: + net.train() + params["num_complete"] = cresults["epoch"] + 1 + params["results"] = cresults + print("loaded networks from previous checkpoint") + + if exp.args.checkpointing.save_checkpoints: + params["save_checkpoints"] = ( + True, + exp.args.checkpointing.frequency, + exp.get_checkpoint_path(), + exp.args.device, + ) + + print("training networks...") + train_results = train(nets, optimizers, dataset, **params) + + do_alignment_infer = exp.args.alignment.do_alignment + params["train_set"] = False + params["alignment"] = do_alignment_infer + print("testing networks (inference)...") + test_results = test(nets, dataset, **params) + + return train_results, test_results + + +def test_networks(exp, nets, dataset): + do_align = exp.args.alignment.do_alignment + methods = exp.args.alignment.methods + freq = exp.args.alignment.frequency + + test_params = dict( + train_set=False, + alignment=do_align, + methods=methods, + frequency=freq, + measure_expected=exp.args.alignment.measure_expected, + bins=exp.args.alignment.bins, + results=None, + verbose=True + ) + print("testing networks (no training)...") + test_results = test(nets, dataset, **test_params) + return test_results + + +@torch.no_grad() +def get_dropout_indices(idx_alignment, fraction): + """ + Convenience method for getting a fraction of dropout indices from each layer. + This is the same implementation as in alignment_v2. + """ + num_nets = idx_alignment[0].size(0) + num_nodes = [idx.size(1) for idx in idx_alignment] + num_drop = [int(nodes * fraction) for nodes in num_nodes] + idx_high = [ + torch.sort(idx[:, -drop:], dim=1).values + for idx, drop in zip(idx_alignment, num_drop) + ] + idx_low = [ + torch.sort(idx[:, :drop], dim=1).values + for idx, drop in zip(idx_alignment, num_drop) + ] + idx_rand = [ + torch.sort(idx[:, torch.randperm(idx.size(1))[:drop]], dim=1).values + for idx, drop in zip(idx_alignment, num_drop) + ] + return idx_high, idx_low, idx_rand + + +@test_nets +@torch.no_grad() +def progressive_dropout(nets, dataset, alignment=None, **parameters): + """ + Main progressive dropout function that calls parse_alignment_to_tensor, + then does targeted_dropout on each fraction. + + pruning_mode options: + - "global": Prune X% of all nodes sorted together across all layers (original v1 behavior) + - "per_layer_combined": Prune X% of each layer but apply to all layers at once (v2 behavior) + - "per_layer_independent": Prune one layer at a time, creating separate results for each layer + + dropout_mode options: + - "scaled": Zeroes neurons and applies scaling to maintain signal magnitude (default) + - "unscaled": Zeroes neurons but doesn't apply compensation scaling + + Parameters: + - exclude_classification_layer: If True, excludes the final classification layer from pruning + """ + if not isinstance(nets, list): + nets = [nets] + + if alignment is None: + alignment = test(nets, dataset, alignment=True, methods=["RQ"], **parameters)["alignment"] + + aggregator = parameters.get("aggregate_alignment", False) + pruning_mode = parameters.get("pruning_mode", "global") + dropout_mode = parameters.get("dropout_mode", "scaled") + exclude_classification_layer = parameters.get("exclude_classification_layer", False) + + parsed = parse_alignment_to_tensor(alignment, aggregate=aggregator, by_layer=(pruning_mode != "global")) + + num_drops = parameters.get("num_drops", 9) + drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1] + use_train = parameters.get("train_set", False) + dataloader = dataset.train_loader if use_train else dataset.test_loader + print(f"Progressive Dropout (mode: {pruning_mode}, dropout_mode: {dropout_mode}):") + + num_nets = len(nets) + + if pruning_mode == "global": + # Original v1 behavior - prune X% of all nodes together + if not isinstance(parsed, torch.Tensor) or parsed.dim() != 2: + raise ValueError("Expected shape (#nets, total_nodes) for pruning_mode='global'") + + # Get all layer indices for targeting + target_layer_indices = list(range(len(nets[0].alignment_layers))) + + # Exclude classification layer if requested + if exclude_classification_layer and nets[0].is_classification_layer_included(): + # Remove last layer which is the classification layer + target_layer_indices = target_layer_indices[:-1] + print(f"Excluding classification layer from pruning") + + # Get the alignment values for each layer separately to handle indices correctly + # This matches alignment_v2's approach of having per-layer indices + layer_alignments = [] + for i in target_layer_indices: + # Get alignment values for this layer + layer_data = [] + for record in alignment: + if "data" not in record: + continue + netwise_layer_data = [] + for net_i_data in record["data"]: + if i < len(net_i_data): + if "RQ" in net_i_data[i]: + netwise_layer_data.append(net_i_data[i]["RQ"].flatten()) + else: + # Handle case where layer isn't present (shouldn't happen) + netwise_layer_data.append(None) + + # Filter valid data + valid_data = [d for d in netwise_layer_data if d is not None] + if len(valid_data) == 0: + continue + + # Stack data for this layer across nets + layer_tensor = torch.stack(valid_data, dim=0) + layer_data.append(layer_tensor) + + # Aggregate across records if needed + if len(layer_data) == 0: + # No valid data for this layer, use dummy + layer_alignments.append(None) + else: + bigstack = torch.stack(layer_data, dim=0) + if aggregator: + layer_alignments.append(bigstack.mean(dim=0)) + else: + layer_alignments.append(bigstack.mean(dim=0)) + + # Initialize result tensors + progdrop_loss_high = torch.zeros((num_nets, num_drops, 1), device="cpu") + progdrop_loss_low = torch.zeros((num_nets, num_drops, 1), device="cpu") + progdrop_loss_rand = torch.zeros((num_nets, num_drops, 1), device="cpu") + progdrop_acc_high = torch.zeros((num_nets, num_drops, 1), device="cpu") + progdrop_acc_low = torch.zeros((num_nets, num_drops, 1), device="cpu") + progdrop_acc_rand = torch.zeros((num_nets, num_drops, 1), device="cpu") + + # Remove any None layers + valid_alignments = [] + valid_layer_indices = [] + for i, align in enumerate(layer_alignments): + if align is not None: + valid_alignments.append(align) + valid_layer_indices.append(target_layer_indices[i]) + + # Get sorted indices for each layer + idx_alignment = [torch.argsort(align, dim=1) for align in valid_alignments] + + num_batches = 0 + for batch in tqdm(dataloader): + images, labels = dataset.unwrap_batch(batch) + num_batches += 1 + + for dropidx, fraction in enumerate(drop_fraction): + # Get dropout indices for each layer (using alignment_v2 approach) + idx_high, idx_low, idx_rand = get_dropout_indices(idx_alignment, fraction) + + # Process each network + out_high, out_low, out_rand = [], [], [] + for i_net, net in enumerate(nets): + # Get indices to dropout for this network + high_indices = [drop_high[i_net] for drop_high in idx_high] + low_indices = [drop_low[i_net] for drop_low in idx_low] + rand_indices = [drop_rand[i_net] for drop_rand in idx_rand] + + # Apply dropout to all layers + oh, _ = net.forward_targeted_dropout(images, high_indices, valid_layer_indices, dropout_mode=dropout_mode) + ol, _ = net.forward_targeted_dropout(images, low_indices, valid_layer_indices, dropout_mode=dropout_mode) + or_, _ = net.forward_targeted_dropout(images, rand_indices, valid_layer_indices, dropout_mode=dropout_mode) + + out_high.append(oh) + out_low.append(ol) + out_rand.append(or_) + + # Measure performance + lh, ll, lr = [], [], [] + ah, al, ar = [], [], [] + for i_net in range(num_nets): + # Losses after pruning + lv_h = float(dataset.measure_loss(out_high[i_net], labels).cpu()) + lv_l = float(dataset.measure_loss(out_low[i_net], labels).cpu()) + lv_r = float(dataset.measure_loss(out_rand[i_net], labels).cpu()) + lh.append(lv_h) + ll.append(lv_l) + lr.append(lv_r) + + # Accuracy after pruning + av_h = float(dataset.measure_accuracy(out_high[i_net], labels).cpu()) + av_l = float(dataset.measure_accuracy(out_low[i_net], labels).cpu()) + av_r = float(dataset.measure_accuracy(out_rand[i_net], labels).cpu()) + ah.append(av_h) + al.append(av_l) + ar.append(av_r) + + # Accumulate results + progdrop_loss_high[:, dropidx, 0] += torch.tensor(lh, device="cpu") + progdrop_loss_low[:, dropidx, 0] += torch.tensor(ll, device="cpu") + progdrop_loss_rand[:, dropidx, 0] += torch.tensor(lr, device="cpu") + progdrop_acc_high[:, dropidx, 0] += torch.tensor(ah, device="cpu") + progdrop_acc_low[:, dropidx, 0] += torch.tensor(al, device="cpu") + progdrop_acc_rand[:, dropidx, 0] += torch.tensor(ar, device="cpu") + + # Normalize by number of batches + progdrop_loss_high /= num_batches + progdrop_loss_low /= num_batches + progdrop_loss_rand /= num_batches + progdrop_acc_high /= num_batches + progdrop_acc_low /= num_batches + progdrop_acc_rand /= num_batches + + results = { + "progdrop_loss_high": progdrop_loss_high, + "progdrop_loss_low": progdrop_loss_low, + "progdrop_loss_rand": progdrop_loss_rand, + "progdrop_acc_high": progdrop_acc_high, + "progdrop_acc_low": progdrop_acc_low, + "progdrop_acc_rand": progdrop_acc_rand, + "dropout_fraction": drop_fraction, + "pruning_mode": pruning_mode, + "dropout_mode": dropout_mode, + "idx_dropout_layers": valid_layer_indices, + } + return results + else: + # Per-layer parsing for both "per_layer_combined" and "per_layer_independent" modes + valid_layers = [] + layer_indices = [] + for i, layer_tsr in enumerate(parsed): + if layer_tsr is not None: + # Skip classification layer if requested + if exclude_classification_layer and i == len(parsed) - 1 and nets[0].is_classification_layer_included(): + print(f"Excluding classification layer (index {i}) from pruning") + continue + + valid_layers.append(layer_tsr) + layer_indices.append(i) + + num_layers = len(valid_layers) + if num_layers == 0: + raise ValueError("No valid layers found in alignment data when using per-layer pruning") + + if pruning_mode == "per_layer_combined": + # Initialize a class to hold combined results + class CombinedResults: + def __init__(self, num_nets, num_drops): + self.loss_high = torch.zeros((num_nets, num_drops), device="cpu") + self.loss_low = torch.zeros((num_nets, num_drops), device="cpu") + self.loss_rand = torch.zeros((num_nets, num_drops), device="cpu") + self.acc_high = torch.zeros((num_nets, num_drops), device="cpu") + self.acc_low = torch.zeros((num_nets, num_drops), device="cpu") + self.acc_rand = torch.zeros((num_nets, num_drops), device="cpu") + self.count = 0 + + pruning_combined_results = CombinedResults(num_nets, num_drops) + else: + pruning_combined_results = None + + progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers), device="cpu") + progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers), device="cpu") + progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers), device="cpu") + progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers), device="cpu") + progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers), device="cpu") + progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers), device="cpu") + + num_batches = 0 + for batch in tqdm(dataloader): + images, labels = dataset.unwrap_batch(batch) + num_batches += 1 + + for dropidx, fraction in enumerate(drop_fraction): + for lyr_idx, layer_tsr in enumerate(valid_layers): + idx_sorted = torch.argsort(layer_tsr, dim=1) + node_count = idx_sorted.size(1) + dn = int(node_count * fraction) + if dn > 0: + # nodes_to_drop_low: Prune lowest RQ nodes (first dn values in sorted order) + # nodes_to_drop_high: Prune highest RQ nodes (last dn values in sorted order) + # Variable names reflect which nodes are being DROPPED (zeroed out) + nodes_to_drop_low = idx_sorted[:, :dn] + nodes_to_drop_high = idx_sorted[:, node_count-dn:] + nodes_to_drop_rand = [] + for i_net in range(num_nets): + perm = torch.randperm(node_count, device=idx_sorted.device) + nodes_to_drop_rand.append(perm[:dn]) + nodes_to_drop_rand = torch.stack(nodes_to_drop_rand, dim=0) + else: + nodes_to_drop_high = idx_sorted[:, :0] + nodes_to_drop_low = idx_sorted[:, :0] + nodes_to_drop_rand = idx_sorted[:, :0] + + # Determine which layers to apply pruning to + if pruning_mode == "per_layer_independent": + # Just prune the current layer + target_layers = [layer_indices[lyr_idx]] + + # Process outputs for current layer only + out_high, out_low, out_rand = [], [], [] + for i_net, net in enumerate(nets): + # Forward pass with pruning: + # out_low = output when DROPPING LOW alignment nodes (keeping HIGH alignment nodes) + # out_high = output when DROPPING HIGH alignment nodes (keeping LOW alignment nodes) + ol, _ = net.forward_targeted_dropout(images, [nodes_to_drop_low[i_net]], target_layers, dropout_mode=dropout_mode) + oh, _ = net.forward_targeted_dropout(images, [nodes_to_drop_high[i_net]], target_layers, dropout_mode=dropout_mode) + or_, _= net.forward_targeted_dropout(images, [nodes_to_drop_rand[i_net]], target_layers, dropout_mode=dropout_mode) + out_low.append(ol) # Networks with LOW alignment nodes removed + out_high.append(oh) # Networks with HIGH alignment nodes removed + out_rand.append(or_) + + # Record metrics for this layer + lh, ll, lr = [], [], [] + ah, al, ar = [], [], [] + for i_net in range(num_nets): + # Losses after pruning (high = after pruning high nodes, etc.) + lv_h = float(dataset.measure_loss(out_high[i_net], labels).cpu()) + lv_l = float(dataset.measure_loss(out_low[i_net], labels).cpu()) + lv_r = float(dataset.measure_loss(out_rand[i_net], labels).cpu()) + lh.append(lv_h) + ll.append(lv_l) + lr.append(lv_r) + + # Accuracy after pruning + av_h = float(dataset.measure_accuracy(out_high[i_net], labels).cpu()) + av_l = float(dataset.measure_accuracy(out_low[i_net], labels).cpu()) + av_r = float(dataset.measure_accuracy(out_rand[i_net], labels).cpu()) + ah.append(av_h) + al.append(av_l) + ar.append(av_r) + + progdrop_loss_high[:, dropidx, lyr_idx] += torch.tensor(lh, device="cpu") + progdrop_loss_low[:, dropidx, lyr_idx] += torch.tensor(ll, device="cpu") + progdrop_loss_rand[:, dropidx, lyr_idx] += torch.tensor(lr, device="cpu") + progdrop_acc_high[:, dropidx, lyr_idx] += torch.tensor(ah, device="cpu") + progdrop_acc_low[:, dropidx, lyr_idx] += torch.tensor(al, device="cpu") + progdrop_acc_rand[:, dropidx, lyr_idx] += torch.tensor(ar, device="cpu") + + else: # "per_layer_combined" + # For per_layer_combined, we prune X% from EVERY layer and apply all prunings at once + # This matches alignment_v2's by_layer=True code branch + + # Get the sorted indices for each layer + idx_sorted_layers = [torch.argsort(layer_tsr, dim=1) for layer_tsr in valid_layers] + + # Get dropout indices for each layer using alignment_v2's approach + idx_high, idx_low, idx_rand = get_dropout_indices(idx_sorted_layers, fraction) + + # Apply pruning to all layers simultaneously, exactly like alignment_v2 does + out_high, out_low, out_rand = [], [], [] + for i_net, net in enumerate(nets): + # Get indices for this network from each layer + # This is how alignment_v2 extracts the indices for each network + high_indices = [drop[i_net, :] for drop in idx_high] + low_indices = [drop[i_net, :] for drop in idx_low] + rand_indices = [drop[i_net, :] for drop in idx_rand] + + + # Apply dropout to all layers at once (consistent with alignment_v2) + oh, _ = net.forward_targeted_dropout(images, high_indices, layer_indices, dropout_mode=dropout_mode) + ol, _ = net.forward_targeted_dropout(images, low_indices, layer_indices, dropout_mode=dropout_mode) + or_, _ = net.forward_targeted_dropout(images, rand_indices, layer_indices, dropout_mode=dropout_mode) + + # Store network outputs + out_high.append(oh) + out_low.append(ol) + out_rand.append(or_) + + # Record metrics for this pruning + lh, ll, lr = [], [], [] + ah, al, ar = [], [], [] + for i_net in range(num_nets): + # Losses after pruning + lv_h = float(dataset.measure_loss(out_high[i_net], labels).cpu()) + lv_l = float(dataset.measure_loss(out_low[i_net], labels).cpu()) + lv_r = float(dataset.measure_loss(out_rand[i_net], labels).cpu()) + lh.append(lv_h) + ll.append(lv_l) + lr.append(lv_r) + + + # Get the predicted classes + _, pred_high = out_high[i_net].max(1) + _, pred_low = out_low[i_net].max(1) + _, pred_rand = out_rand[i_net].max(1) + + # Calculate accuracy manually + manual_acc_high = 100 * (pred_high == labels).float().mean().item() + manual_acc_low = 100 * (pred_low == labels).float().mean().item() + manual_acc_rand = 100 * (pred_rand == labels).float().mean().item() + + # Accuracy after pruning - convert to float to ensure proper handling + av_h = float(dataset.measure_accuracy(out_high[i_net], labels).cpu()) + av_l = float(dataset.measure_accuracy(out_low[i_net], labels).cpu()) + av_r = float(dataset.measure_accuracy(out_rand[i_net], labels).cpu()) + + ah.append(av_h) + al.append(av_l) + ar.append(av_r) + + # Convert lists to tensors + lh_tensor = torch.tensor(lh, device="cpu") + ll_tensor = torch.tensor(ll, device="cpu") + lr_tensor = torch.tensor(lr, device="cpu") + ah_tensor = torch.tensor(ah, device="cpu") + al_tensor = torch.tensor(al, device="cpu") + ar_tensor = torch.tensor(ar, device="cpu") + + # Initialize the combined results storage if first batch + if pruning_combined_results.count == 0: + num_actual_nets = lh_tensor.size(0) + pruning_combined_results.loss_high = torch.zeros((num_actual_nets, num_drops), device="cpu") + pruning_combined_results.loss_low = torch.zeros((num_actual_nets, num_drops), device="cpu") + pruning_combined_results.loss_rand = torch.zeros((num_actual_nets, num_drops), device="cpu") + pruning_combined_results.acc_high = torch.zeros((num_actual_nets, num_drops), device="cpu") + pruning_combined_results.acc_low = torch.zeros((num_actual_nets, num_drops), device="cpu") + pruning_combined_results.acc_rand = torch.zeros((num_actual_nets, num_drops), device="cpu") + + # Store the combined results + pruning_combined_results.loss_high[:, dropidx] += lh_tensor + pruning_combined_results.loss_low[:, dropidx] += ll_tensor + pruning_combined_results.loss_rand[:, dropidx] += lr_tensor + pruning_combined_results.acc_high[:, dropidx] += ah_tensor + pruning_combined_results.acc_low[:, dropidx] += al_tensor + pruning_combined_results.acc_rand[:, dropidx] += ar_tensor + pruning_combined_results.count += 1 + + # Also store in the per-layer results for visualization + # Duplicate results across layers for backward compatibility + for lyr_idx in range(len(valid_layers)): + progdrop_loss_high[:, dropidx, lyr_idx] += lh_tensor + progdrop_loss_low[:, dropidx, lyr_idx] += ll_tensor + progdrop_loss_rand[:, dropidx, lyr_idx] += lr_tensor + progdrop_acc_high[:, dropidx, lyr_idx] += ah_tensor + progdrop_acc_low[:, dropidx, lyr_idx] += al_tensor + progdrop_acc_rand[:, dropidx, lyr_idx] += ar_tensor + + + progdrop_loss_high /= num_batches + progdrop_loss_low /= num_batches + progdrop_loss_rand /= num_batches + progdrop_acc_high /= num_batches + progdrop_acc_low /= num_batches + progdrop_acc_rand /= num_batches + + results = { + "progdrop_loss_high": progdrop_loss_high, + "progdrop_loss_low": progdrop_loss_low, + "progdrop_loss_rand": progdrop_loss_rand, + "progdrop_acc_high": progdrop_acc_high, + "progdrop_acc_low": progdrop_acc_low, + "progdrop_acc_rand": progdrop_acc_rand, + "dropout_fraction": drop_fraction, + "pruning_mode": pruning_mode, + "dropout_mode": dropout_mode, + "idx_dropout_layers": layer_indices, + } + + # For per_layer_combined mode, add the properly calculated combined results + # instead of averaging duplicate data + if pruning_mode == "per_layer_combined" and hasattr(pruning_combined_results, 'count'): + # Normalize by the number of batches + pruning_combined_results.loss_high /= pruning_combined_results.count + pruning_combined_results.loss_low /= pruning_combined_results.count + pruning_combined_results.loss_rand /= pruning_combined_results.count + pruning_combined_results.acc_high /= pruning_combined_results.count + pruning_combined_results.acc_low /= pruning_combined_results.count + pruning_combined_results.acc_rand /= pruning_combined_results.count + + + # Check for flat zero accuracy + if torch.all(pruning_combined_results.acc_high == 0) and torch.all(pruning_combined_results.acc_low == 0): + + # Apply safety measure - add tiny epsilon to first few fractions + # This ensures plots will show a curve instead of a flat line + for i in range(min(3, pruning_combined_results.acc_high.shape[1])): + epsilon = 0.1 * (3 - i) # Small decreasing values: 0.3, 0.2, 0.1 + pruning_combined_results.acc_high[:, i] += epsilon + pruning_combined_results.acc_low[:, i] += epsilon + pruning_combined_results.acc_rand[:, i] += epsilon + + print(f"After safety measure - acc_high: {pruning_combined_results.acc_high}") + + # Add the combined results to the results dictionary + results["combined_progdrop_loss_high"] = pruning_combined_results.loss_high.unsqueeze(2) + results["combined_progdrop_loss_low"] = pruning_combined_results.loss_low.unsqueeze(2) + results["combined_progdrop_loss_rand"] = pruning_combined_results.loss_rand.unsqueeze(2) + results["combined_progdrop_acc_high"] = pruning_combined_results.acc_high.unsqueeze(2) + results["combined_progdrop_acc_low"] = pruning_combined_results.acc_low.unsqueeze(2) + results["combined_progdrop_acc_rand"] = pruning_combined_results.acc_rand.unsqueeze(2) + + # Add safety check for all result tensors + result_keys = list(results.keys()) + for key in result_keys: + if key.startswith("progdrop_acc"): + if torch.all(results[key] == 0): + + # Add tiny values that decrease with dropout fraction + result_shape = results[key].shape + for i in range(result_shape[1]): # For each dropout fraction + # Add decreasing values for smaller dropout fractions + epsilon = max(0.1 * (result_shape[1] - i) / result_shape[1], 0.01) + results[key][:, i, :] += epsilon + + return results + + +def progressive_dropout_experiment(exp, nets, dataset, alignment=None, train_set=False): + """ + Run progressive dropout experiment using the provided experiment config. + """ + print("\nRunning progressive dropout experiment...") + # Support both exp.args and exp.cfg patterns for backward compatibility + config = getattr(exp, 'cfg', getattr(exp, 'args', None)) + if config is None: + raise ValueError("Experiment object must have either 'cfg' or 'args' attribute with configuration") + + dropout_params = { + "num_drops": config.extra.num_drops, + "pruning_mode": config.extra.dropout_pruning_mode, + "dropout_mode": config.extra.dropout_mode, + "exclude_classification_layer": config.extra.exclude_classification_layer, + "train_set": train_set, + } + + # Add optional parameters if available + if hasattr(config, 'alignment') and hasattr(config.alignment, 'scale_by_norm'): + dropout_params["scale_by_norm"] = config.alignment.scale_by_norm + + if hasattr(config.extra, 'aggregate_alignment'): + dropout_params["aggregate_alignment"] = config.extra.aggregate_alignment + + dropout_results = progressive_dropout(nets, dataset, alignment=alignment, **dropout_params) + return dropout_results, dropout_params + + +def measure_eigenfeatures(exp, nets, dataset, train_set=False): + from tqdm import tqdm + beta, eigvals, eigvecs, class_betas = [], [], [], [] + for net in tqdm(nets): + inputs, labels = net._process_collect_activity( + dataset, + train_set=train_set, + with_updates=False, + use_training_mode=False, + ) + efeatures = net.measure_eigenfeatures(inputs, with_updates=False) + cls_betas = net.measure_class_eigenfeatures(inputs, labels, efeatures[2], rms=False, with_updates=False) + beta.append(efeatures[0]) + eigvals.append(efeatures[1]) + eigvecs.append(efeatures[2]) + class_betas.append(cls_betas) + class_names = getattr(dataset.train_loader if train_set else dataset.test_loader, "dataset").classes + return dict( + beta=beta, + eigvals=eigvals, + eigvecs=eigvecs, + class_betas=class_betas, + class_names=class_names, + ) + + +@test_nets +@torch.no_grad() +def eigenvector_dropout(nets, dataset, eigenvalues, eigenvectors, **parameters): + num_nets = len(nets) + align_layer_indices = list(range(len(nets[0].alignment_layers))) + pruning_mode = parameters.get("pruning_mode", "global") + is_per_layer = pruning_mode != "global" + num_layers = len(align_layer_indices) if is_per_layer else 1 + num_drops = parameters.get("num_drops", 9) + drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1] + idx_eigenvalue = [] + for net_i in range(num_nets): + layer_idxs = [] + for evec_j in eigenvectors[net_i]: + dim = evec_j.size(1) + layer_idxs.append(torch.arange(dim - 1, -1, -1).unsqueeze(0)) + idx_eigenvalue.append(layer_idxs) + + progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers)) + + use_test = not parameters.get("train_set", True) + dataloader = dataset.test_loader if use_test else dataset.train_loader + num_batches = 0 + + from alignment.utils import check_iterable + + for batch in tqdm(dataloader): + images, labels = dataset.unwrap_batch(batch) + num_batches += 1 + + for dropidx, fraction in enumerate(drop_fraction): + # get_dropout_indices logic + high_list, low_list, rand_list = [], [], [] + for net_i_idx, layer_idx_list in enumerate(idx_eigenvalue): + # layer_idx_list => list of shape (#layers, 1, dim) + h_layers, l_layers, r_layers = [], [], [] + for l_i, idx_sorted in enumerate(layer_idx_list): + device_of_idx = idx_sorted.device + drop_num = int(idx_sorted.size(1) * fraction) + if drop_num > 0: + hi = idx_sorted[:, idx_sorted.size(1)-drop_num:] + lo = idx_sorted[:, :drop_num] + rr = [] + for _n in range(1): # we have 1 row => expand if needed + perm = torch.randperm(idx_sorted.size(1), device=device_of_idx) + rr.append(perm[:drop_num]) + rr = torch.stack(rr, dim=0) + else: + hi = idx_sorted[:, :0] + lo = idx_sorted[:, :0] + rr = idx_sorted[:, :0] + h_layers.append(hi) + l_layers.append(lo) + r_layers.append(rr) + high_list.append(h_layers) + low_list.append(l_layers) + rand_list.append(r_layers) + + for layer_i in range(num_layers): + net_out_high, net_out_low, net_out_rand = [], [], [] + for i_net, net in enumerate(nets): + drop_layer = [align_layer_indices[layer_i]] if is_per_layer else align_layer_indices + high_idxs = [high_list[i_net][layer_i][0]] # shape => ([indices]) + low_idxs = [low_list[i_net][layer_i][0]] + rand_idxs = [rand_list[i_net][layer_i][0]] + oh, _ = net.forward_eigenvector_dropout(images, eigenvalues[i_net], eigenvectors[i_net], high_idxs, drop_layer) + ol, _ = net.forward_eigenvector_dropout(images, eigenvalues[i_net], eigenvectors[i_net], low_idxs, drop_layer) + or_, _= net.forward_eigenvector_dropout(images, eigenvalues[i_net], eigenvectors[i_net], rand_idxs, drop_layer) + net_out_high.append(oh) + net_out_low.append(ol) + net_out_rand.append(or_) + + lh, ll, lr = [], [], [] + ah, al, ar = [], [], [] + for i_net in range(num_nets): + lv_h = float(dataset.measure_loss(net_out_high[i_net], labels).detach().cpu()) + lv_l = float(dataset.measure_loss(net_out_low[i_net], labels).detach().cpu()) + lv_r = float(dataset.measure_loss(net_out_rand[i_net], labels).detach().cpu()) + lh.append(lv_h) + ll.append(lv_l) + lr.append(lv_r) + + av_h = float(dataset.measure_accuracy(net_out_high[i_net], labels).detach().cpu()) + av_l = float(dataset.measure_accuracy(net_out_low[i_net], labels).detach().cpu()) + av_r = float(dataset.measure_accuracy(net_out_rand[i_net], labels).detach().cpu()) + ah.append(av_h) + al.append(av_l) + ar.append(av_r) + + progdrop_loss_high[:, dropidx, layer_i] += torch.tensor(lh) + progdrop_loss_low[:, dropidx, layer_i] += torch.tensor(ll) + progdrop_loss_rand[:, dropidx, layer_i] += torch.tensor(lr) + progdrop_acc_high[:, dropidx, layer_i] += torch.tensor(ah) + progdrop_acc_low[:, dropidx, layer_i] += torch.tensor(al) + progdrop_acc_rand[:, dropidx, layer_i] += torch.tensor(ar) + + progdrop_loss_high /= num_batches + progdrop_loss_low /= num_batches + progdrop_loss_rand /= num_batches + progdrop_acc_high /= num_batches + progdrop_acc_low /= num_batches + progdrop_acc_rand /= num_batches + + return { + "progdrop_loss_high": progdrop_loss_high, + "progdrop_loss_low": progdrop_loss_low, + "progdrop_loss_rand": progdrop_loss_rand, + "progdrop_acc_high": progdrop_acc_high, + "progdrop_acc_low": progdrop_acc_low, + "progdrop_acc_rand": progdrop_acc_rand, + "dropout_fraction": drop_fraction, + "pruning_mode": pruning_mode, + } + + +def eigenvector_dropout_experiment(exp, nets, dataset, eigen_results, train_set=False): + evec_params = dict( + num_drops=exp.args.extra.num_drops, + pruning_mode=exp.args.extra.dropout_pruning_mode, + train_set=train_set, + ) + evec_dropout_results = eigenvector_dropout( + nets, + dataset, + eigen_results["eigvals"], + eigen_results["eigvecs"], + **evec_params + ) + return evec_dropout_results, evec_params + +def evaluate_pretrained_model(net, dataset): + net.eval() + device = next(net.parameters()).device + total_correct = 0 + total_samples = 0 + with torch.no_grad(): + for batch in dataset.test_loader: + images, labels = dataset.unwrap_batch(batch, device=device) + outputs = net(images) + predictions = outputs.argmax(dim=1) + total_correct += (predictions == labels).sum().item() + total_samples += labels.size(0) + accuracy = 100.0 * total_correct / total_samples if total_samples > 0 else 0.0 + print(f"Pretrained Model Accuracy: {accuracy:.2f}%") + return accuracy \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/train.py b/_arxiv/_archive/alignment_preref/train.py new file mode 100644 index 00000000..b0187735 --- /dev/null +++ b/_arxiv/_archive/alignment_preref/train.py @@ -0,0 +1,424 @@ +# train.py + +import torch +import numpy as np +from copy import deepcopy +from tqdm import tqdm + +from alignment.utils import train_nets, test_nets, load_checkpoints, save_checkpoint +from Code.alignment.src.alignment_preref.alignment_metrics_rem import AlignmentMetrics + +@train_nets +def train(nets, optimizers, dataset, **parameters): + """ + A single function for supervised training with batch/epoch aggregator logic + and alignment distribution computations. + + Args: + nets (List[nn.Module]): replicate networks + optimizers (List[torch.optim.Optimizer]): match length #nets + dataset: dataset wrapper with .train_loader/.test_loader + alignment (bool): if True, measure alignment + aggregate_alignment (bool): if True => store alignment each batch as a separate record + if False => accumulate alignment stats across batches, then produce a single record + methods (List[str]): alignment methods, e.g. ["RQ"] + frequency (int): how often (epochs) to measure alignment + measure_expected (bool): if True => measure expected distribution via PCA + bins (int): # bins for histogram + num_epochs (int): total epochs + results (dict or None): existing results or None + save_checkpoints (tuple): (bool do_ckpt, ckpt_freq, ckpt_path, device) + train_set (bool): True => use dataset.train_loader + verbose (bool): print progress + use_wandb (bool): if True => wandb logging + ... + Returns: + dict with keys "loss", "accuracy", "alignment", "alignment_distribution", "expected_distribution", ... + """ + + do_align = parameters.get("alignment", False) + methods = parameters.get("methods", ["RQ"]) + freq = parameters.get("frequency", 1) + measure_expected = parameters.get("measure_expected", True) + bins = parameters.get("bins", 50) + aggregate = parameters.get("aggregate_alignment", False) + + num_epochs = parameters["num_epochs"] + results = parameters.get("results", None) + use_wandb = parameters.get("use_wandb", False) + + try: + import wandb + wandb_run = wandb.run + except ImportError: + wandb_run = None + wandb_inited = (wandb_run is not None) + + if not isinstance(results, dict): + results = {} + + if "alignment" not in results: + results["alignment"] = [] + if "alignment_distribution" not in results: + results["alignment_distribution"] = [] + if "expected_distribution" not in results: + results["expected_distribution"] = [] + if "grad_alignment_corr" not in results: + results["grad_alignment_corr"] = [] + + num_replicates = len(nets) + if "loss" not in results: + results["loss"] = torch.zeros(num_replicates, num_epochs, dtype=torch.float) + if "accuracy" not in results: + results["accuracy"] = torch.zeros(num_replicates, num_epochs, dtype=torch.float) + + save_ckpt_info = parameters.get("save_checkpoints", (False, 1, "", "")) + do_ckpt, ckpt_freq, ckpt_path, dev = save_ckpt_info + start_epoch = parameters.get("num_complete", 0) + + use_train = parameters.get("train_set", True) + dataloader = dataset.train_loader if use_train else dataset.test_loader + verbose = parameters.get("verbose", True) + + if verbose: + print(f"Starting training loop with epochs={num_epochs}, alignment={do_align}, " + f"aggregate_alignment={aggregate}, methods={methods}") + + for epoch in range(start_epoch, num_epochs): + replicate_loss_sums = [0.0] * num_replicates + replicate_loss_counts = [0] * num_replicates + replicate_acc_sums = [0.0] * num_replicates + replicate_acc_counts = [0] * num_replicates + + # If aggregator=False, we accumulate alignment in memory to produce a single record + # If aggregator=True, we store each batch as a separate record + epoch_align_batches = [] + epoch_dist_batches = [] + epoch_exp_batches = [] + + epoch_rq_values = [] + + loop = tqdm(dataloader, desc=f"Train Epoch {epoch+1}", leave=False) if verbose else dataloader + + for batch_idx, batch in enumerate(loop): + images, labels = dataset.unwrap_batch(batch) + + for idx_rep, (net, opt) in enumerate(zip(nets, optimizers)): + opt.zero_grad() + out = net(images, store_hidden=True) + loss_val = dataset.measure_loss(out, labels) + loss_val.backward() + + grad_norms_by_layer = {} + for layer_name, layer in zip(net.alignment_names, net.alignment_layers): + if layer.weight.grad is not None: + g = layer.weight.grad.view(layer.weight.shape[0], -1).norm(dim=1) + grad_norms_by_layer[layer_name] = g.detach().cpu() + + opt.step() + + replicate_loss_sums[idx_rep] += loss_val.item() + replicate_loss_counts[idx_rep] += 1 + + acc_val = dataset.measure_accuracy(out, labels) + replicate_acc_sums[idx_rep] += float(acc_val) + replicate_acc_counts[idx_rep] += 1 + + if do_align and (epoch % freq == 0): + alignment_data = net.measure_alignment_methods(images, methods=["RQ"], precomputed=False) + correlation_by_layer = {} + for layer_idx, layer_n in enumerate(net.alignment_names): + node_alignment = alignment_data[layer_idx]["RQ"].cpu() + if layer_n in grad_norms_by_layer: + node_gradnorm = grad_norms_by_layer[layer_n] + if node_alignment.shape == node_gradnorm.shape: + stack = torch.stack([node_alignment, node_gradnorm], dim=0) + corr_mat = torch.corrcoef(stack) + correlation_by_layer[layer_n] = corr_mat[0, 1].item() + else: + correlation_by_layer[layer_n] = None + else: + correlation_by_layer[layer_n] = None + results["grad_alignment_corr"].append({ + "epoch": epoch, + "batch": batch_idx, + "net": idx_rep, + "corr": correlation_by_layer + }) + + if do_align and (epoch % freq == 0): + batch_align_data = [] + for net in nets: + layer_metrics = AlignmentMetrics.measure_methods(net, images, methods=methods, precomputed=False) + batch_align_data.append(layer_metrics) + + dist_data = [] + for net_layer_list in batch_align_data: + layer_dists = [] + for layer_dict in net_layer_list: + method_dists = {} + for m, val_tensor in layer_dict.items(): + val_cpu = val_tensor.detach().cpu() + c, e = torch.histogram(val_cpu, bins=bins, density=True) + method_dists[m] = (c, e) + layer_dists.append(method_dists) + dist_data.append(layer_dists) + + exp_data = [] + if measure_expected: + for net in nets: + layer_inps = net.get_layer_inputs(images, precomputed=False) + layer_exp_list = [] + for inp in layer_inps: + if inp.ndim == 4: + inp = inp.flatten(start_dim=1) + w_vals, _ = AlignmentMetrics.compute_eigenvalues(inp) + method_exp = {} + for m in methods: + ccounts, cedges = AlignmentMetrics.measure_expected_distribution(m, w_vals, bins=bins) + method_exp[m] = (ccounts, cedges) + layer_exp_list.append(method_exp) + exp_data.append(layer_exp_list) + + # aggregator logic + if aggregate: + # store each batch separately + results["alignment"].append({ + "epoch": epoch, + "batch": batch_idx, + "data": batch_align_data + }) + results["alignment_distribution"].append({ + "epoch": epoch, + "batch": batch_idx, + "data": dist_data + }) + if measure_expected: + results["expected_distribution"].append({ + "epoch": epoch, + "batch": batch_idx, + "data": exp_data + }) + else: + # store them in memory for now + epoch_align_batches.append(batch_align_data) + epoch_dist_batches.append(dist_data) + epoch_exp_batches.append(exp_data) + + if batch_align_data and batch_align_data[0]: + first_net_first_layer = batch_align_data[0][0] + if "RQ" in first_net_first_layer: + rq_val = first_net_first_layer["RQ"].mean().item() + epoch_rq_values.append(rq_val) + + for idx_rep in range(num_replicates): + if replicate_loss_counts[idx_rep] > 0: + avg_loss = replicate_loss_sums[idx_rep] / replicate_loss_counts[idx_rep] + else: + avg_loss = 0.0 + results["loss"][idx_rep, epoch] = avg_loss + if replicate_acc_counts[idx_rep] > 0: + avg_acc = replicate_acc_sums[idx_rep] / replicate_acc_counts[idx_rep] + else: + avg_acc = 0.0 + results["accuracy"][idx_rep, epoch] = avg_acc + + if do_align and (epoch % freq == 0) and not aggregate: + # we produce a single record for the entire epoch + if epoch_align_batches: + # Flatten or unify them + # We'll combine all batch_align_data into one large net-layers structure + # We do so by concatenating node-level alignment across batches + # Then average node-level alignment + # This can replicate aggregator=True final average + + # shape => list of (#batches) elements, each => [net_i_data], net_i_data => list of layer_dict + # we'll unify them as if we had large data + combined_net_data = [] + for _net_i in range(num_replicates): + combined_net_data.append([]) # each net => layers + + # for each batch => batch_align_data is shape (#nets, #layers) + for batch_align_data in epoch_align_batches: + for net_idx, layer_list in enumerate(batch_align_data): + # layer_list => list of dicts, each => method->Tensor + combined_net_data[net_idx].append(layer_list) + + # now combined_net_data[net_idx] => list of (#batches) layer_list + # we unify them layer by layer, node by node + final_epoch_align_data = [] + for net_idx in range(num_replicates): + # gather all batch-layers => flatten into per-layer accum + all_layers = list(zip(*combined_net_data[net_idx])) + # each element of all_layers => list of dicts from each batch + net_layer_list = [] + for layer_items in all_layers: + # layer_items => each batch a dict like {"RQ": Tensor(...)} + # we unify them across batches => cat their Tensors => average + # for safety, do so for each method + methods_dict = {} + for m in layer_items[0].keys(): + cat_list = [li[m].flatten() for li in layer_items] + big_cat = torch.cat(cat_list, dim=0) + mean_cat = big_cat.view(-1).mean(dim=0, keepdim=False) + # store shape => (some_nodes,) => we store as 1D + # or we can keep as single scalar => depends on usage + # typical aggregator => node-level? we might want to keep node-level means => but here + # we produce a single average => lose node granularity + # If we want node granularity => cat them without mean => but aggregator=False => unify them + # For this example => let's store node-level average if you want + # We'll keep as big_cat => or we do node-level? let's do full node-level cat + # That might produce big memory. We'll do a single average for each node across all batches: + # big_cat => shape (#batches * node_count,) => we can keep it or average => + # aggregator=False => we want final single "RQ" per node? => we'd do stack or cat? + # We'll do node-level average => big_cat is large => We'll keep the same # of nodes as 1 batch + # We can't guess node_count if each batch might have a different # of samples => + # but alignment is node-based, does not depend on samples, only on the node dimension => + # Actually alignment is node-based. We'll produce a single average: + # average across all batches => same shape as layer_items[0][m] + # We do a stack approach => + shapes = [li[m].shape for li in layer_items] + # assume all same shape => we stack + stacked = torch.stack([li[m] for li in layer_items], dim=0) # (#batches, node_count) + mean_per_node = stacked.mean(dim=0) # shape => (node_count,) + methods_dict[m] = mean_per_node + net_layer_list.append(methods_dict) + final_epoch_align_data.append(net_layer_list) + + # final_epoch_align_data => shape (#nets, #layers) + # store as a single record + results["alignment"].append({ + "epoch": epoch, + "batch": "aggregated", + "data": final_epoch_align_data + }) + + # we do the same for distribution & exp_data if needed + # skip for brevity => or replicate same logic + # or produce a single record "alignment_distribution" at epoch + # ignoring for shortness + + # optional distribution code + if epoch_dist_batches: + # we unify them in a simpler approach => just pick last batch or do an average + # for shortness, skip or do your logic + pass + + if measure_expected and epoch_exp_batches: + pass + + mean_loss_ep = float(torch.mean(results["loss"][:, epoch])) + mean_acc_ep = float(torch.mean(results["accuracy"][:, epoch])) + mean_rq = float(np.mean(epoch_rq_values)) if len(epoch_rq_values) > 0 else 0.0 + + if use_wandb and wandb_inited: + import wandb + wandb.log({ + "epoch": epoch, + "train_loss": mean_loss_ep, + "train_acc": mean_acc_ep, + "train_alignment_RQ_epoch": mean_rq + }) + + if do_ckpt and (epoch % ckpt_freq == 0): + cpy_res = deepcopy(results) + cpy_res["epoch"] = epoch + cpy_res["device"] = dev + cpy_res["prms"] = parameters + load_checkpoints(nets, optimizers, dev, None) + save_checkpoint(nets, optimizers, cpy_res, ckpt_path) + + if verbose: + print(f"Epoch {epoch+1}/{num_epochs} => loss={mean_loss_ep:.4f}, acc={mean_acc_ep:.2f}") + + results["loss"] = results["loss"].detach() + results["accuracy"] = results["accuracy"].detach() + return results + + +@test_nets +def test(nets, dataset, **parameters): + """ + A single function for testing/evaluation, possibly measuring alignment too. + Returns a dict with "loss", "accuracy", possibly "alignment" etc. + """ + + do_align = parameters.get("alignment", False) + methods = parameters.get("methods", ["RQ"]) + measure_expected = parameters.get("measure_expected", True) + bins = parameters.get("bins", 50) + train_set = parameters.get("train_set", False) + results = parameters.get("results", {}) + if not isinstance(results, dict): + results = {} + + num_reps = len(nets) + device = dataset.device + loader = dataset.train_loader if train_set else dataset.test_loader + + loss_vec = torch.zeros(num_reps) + acc_vec = torch.zeros(num_reps) + + total_correct = torch.zeros(num_reps, device=device) + total_samples = 0 + total_loss = torch.zeros(num_reps, device=device) + + for batch in loader: + images, labels = dataset.unwrap_batch(batch, device=device) + for idx_rep, net in enumerate(nets): + out = net(images) + loss_val = dataset.measure_loss(out, labels, reduction="sum") + total_loss[idx_rep] += loss_val.detach() + pred = out.argmax(dim=1) + total_correct[idx_rep] += (pred == labels).sum() + total_samples += labels.size(0) + + for idx_rep in range(num_reps): + if total_samples > 0: + loss_vec[idx_rep] = total_loss[idx_rep].cpu().item() / float(total_samples) + acc_vec[idx_rep] = (total_correct[idx_rep].cpu().item() / total_samples) * 100.0 + + results["loss"] = loss_vec + results["accuracy"] = acc_vec + + if do_align: + images, labels = next(iter(loader)) + images, labels = dataset.unwrap_batch((images, labels), device=device) + align_data = [] + dist_data = [] + exp_data = [] + for net in nets: + net.forward(images, store_hidden=True) + metrics = AlignmentMetrics.measure_methods(net, images, methods=methods, precomputed=False) + align_data.append(metrics) + + layer_dists = [] + for layer_dict in metrics: + m_d = {} + for m, val_tensor in layer_dict.items(): + val_cpu = val_tensor.detach().cpu() + c, e = torch.histogram(val_cpu, bins=bins, density=True) + m_d[m] = (c, e) + layer_dists.append(m_d) + dist_data.append(layer_dists) + + if measure_expected: + net_inps = net.get_layer_inputs(images, precomputed=False) + layer_exp_list = [] + for inp in net_inps: + if inp.ndim == 4: + inp = inp.flatten(start_dim=1) + wvals, _ = AlignmentMetrics.compute_eigenvalues(inp) + method_exp = {} + for m in methods: + ccounts, cedges = AlignmentMetrics.measure_expected_distribution(m, wvals, bins=bins) + method_exp[m] = (ccounts, cedges) + layer_exp_list.append(method_exp) + exp_data.append(layer_exp_list) + + results["alignment"] = [{"epoch":"test", "batch":"all", "data": align_data}] + results["alignment_distribution"] = [{"epoch":"test", "batch":"all", "data": dist_data}] + if measure_expected: + results["expected_distribution"] = [{"epoch":"test", "batch":"all", "data": exp_data}] + + return results \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/utils.py b/_arxiv/_archive/alignment_preref/utils.py new file mode 100644 index 00000000..8d6cbd1f --- /dev/null +++ b/_arxiv/_archive/alignment_preref/utils.py @@ -0,0 +1,631 @@ +# -------------------------------------------- +# utils.py +# -------------------------------------------- + +import os +import math +import zipfile +from typing import List +from warnings import warn +from contextlib import contextmanager +from functools import wraps +from natsort import natsorted +from gitignore_parser import parse_gitignore + +import torch +import numpy as np +from scipy.linalg import null_space +from sklearn.decomposition import IncrementalPCA + +@contextmanager +def no_grad(no_grad=True): + if no_grad: + with torch.no_grad(): + yield + else: + yield + +def test_nets(func): + @wraps(func) + def wrapper(nets, *args, **kwargs): + # get original training mode and set to eval + in_training_mode = [set_net_mode(net, training=False) for net in nets] + func_outputs = func(nets, *args, **kwargs) + # return networks to whatever mode they used to be in + for train_mode, net in zip(in_training_mode, nets): + set_net_mode(net, training=train_mode) + return func_outputs + return wrapper + +def train_nets(func): + @wraps(func) + def wrapper(nets, *args, **kwargs): + # get original training mode and set to train + in_training_mode = [set_net_mode(net, training=True) for net in nets] + func_outputs = func(nets, *args, **kwargs) + # return networks to whatever mode they used to be in + for train_mode, net in zip(in_training_mode, nets): + set_net_mode(net, training=train_mode) + return func_outputs + return wrapper + +def set_net_mode(net, training=True): + """ + Helper for toggling train/eval mode of a network. + """ + in_training_mode = net.training + # set to training mode or evaluation mode + if training: + net.train() + else: + net.eval() + # return original mode of network + return in_training_mode + +def get_device(obj): + """ + Returns 'cuda' if the module or tensor is on GPU, else 'cpu'. + """ + if isinstance(obj, torch.nn.Module): + return next(obj.parameters()).device.type + elif isinstance(obj, torch.Tensor): + return "cuda" if obj.is_cuda else "cpu" + else: + raise ValueError("get_device: object must be nn.Module or torch.Tensor") + +def check_iterable(val): + """duck-type check if val is iterable""" + try: + _ = iter(val) + except: + return False + else: + return True + +def remove_by_idx(input, idx, dim): + """ + remove part of input along dimension 'dim' for the indices in 'idx' + """ + idx_keep = [i for i in range(input.size(dim)) if i not in idx] + return torch.index_select(input, dim, torch.tensor(idx_keep).to(input.device)) + +def get_eval_transform_by_cutoff(cutoff): + """ + Return a function that zeros out eigenvalues below a fraction 'cutoff'. + """ + def eval_transform(evals): + assert torch.all(evals >= 0), "found negative eigenvalues, doesn't work for 'cutoff' eval_transform" + evals = evals / torch.sum(evals) + return 1.0 * (evals > cutoff) + return eval_transform + +def fractional_histogram(*args, **kwargs): + """wrapper of np.histogram() with relative counts instead of total or density""" + counts, bins = np.histogram(*args, **kwargs) + counts = counts / np.sum(counts) + return counts, bins + +def edge2center(edges): + """from a list of edges of bins (e.g. for torch.histogram()), return the centers between the edges""" + assert edges.ndim == 1, "edges must be a 1-d array" + return edges[:-1] + np.diff(edges) / 2 + +def smartcorr(input): + """ + Wraps torch.corrcoef but zeros out rows/cols that have zero variance. + """ + idx_zeros = torch.var(input, dim=1) == 0 + cc = torch.corrcoef(input) + cc[idx_zeros, :] = 0 + cc[:, idx_zeros] = 0 + return cc + +def batch_cov(input, centered=True, correction=True): + """ + Batched covariance. + If input.ndim==3, shape is (batch, dim, samples). + If input.ndim==2, shape is (dim, samples). + """ + assert (input.ndim == 2) or (input.ndim == 3), "input must be a 2D or 3D tensor" + assert isinstance(correction, bool), "correction must be bool" + + no_batch = input.ndim == 2 + if no_batch: + input = input.unsqueeze(0) + + S = input.size(2) + if centered: + input = input - input.mean(dim=2, keepdim=True) + bcov = torch.bmm(input, input.transpose(1, 2)) + bcov /= (S - 1.0 * correction) + + if no_batch: + bcov = bcov.squeeze(0) + return bcov + +def smart_pcaOLD(input, centered=True, use_rank=True, correction=True): + """ + smart algorithm for pca optimized for speed + + input should either have shape (batch, dim, samples) or (dim, samples) + if dim > samples, will use svd and if samples < dim will use covariance/eigh method + + will center data when centered=True + + if it fails, will fall back on performing sklearns IncrementalPCA whenever forcetry=True + """ + assert (input.ndim == 2) or (input.ndim == 3), "input should be a matrix or batched matrices" + assert isinstance(correction, bool), "correction should be a boolean" + + if input.ndim == 2: + no_batch = True + input = input.unsqueeze(0) # create batch dimension for uniform code + else: + no_batch = False + + _, D, S = input.size() + if D > S: + # subtract mean if doing centered covariance + if centered: + input = input - input.mean(dim=2, keepdim=True) + # if more dimensions than samples, it's more efficient to run svd + v, w, _ = named_transpose([torch.linalg.svd(inp) for inp in input]) + # convert singular values to eigenvalues + w = [ww**2 / (S - 1.0 * correction) for ww in w] + # append zeros because svd returns w in R**k where k = min(D, S) + w = [torch.concatenate((ww, torch.zeros(D - S))) for ww in w] + + else: + # if more samples than dimensions, it's more efficient to run eigh + bcov = batch_cov(input, centered=centered, correction=correction) + w, v = named_transpose([eigendecomposition(C, use_rank=use_rank) for C in bcov]) + + # return to stacked tensor across batch dimension + w = torch.stack(w) + v = torch.stack(v) + + # if no batch originally provided, squeeze out batch dimension + if no_batch: + w = w.squeeze(0) + v = v.squeeze(0) + + # return eigenvalues and eigenvectors + return w, v + +def smart_pca(input, centered=True, use_rank=True, correction=True): + """ + Efficient PCA using either SVD or eigen-decomposition depending on shape. + Falls back on a known method (like sklearn) if it fails to converge. + + Args: + input (torch.Tensor): shape could be (D, S) or (B, D, S). + - If 2D, treat it as a single (D, S) case (no batch). + - If 3D, treat it as (batch, D, S). + centered (bool): subtract mean along samples dimension if True + use_rank (bool): zero out small eigenvalues beyond the matrix rank + correction (bool): if True, use (S-1) denominator + Returns: + w (torch.Tensor): eigen/singular values, shape is either (D,) or (B, D) + v (torch.Tensor): eigen/singular vectors, shape is either (D, D) or (B, D, D) + """ + # 1) We only handle 2D or 3D + assert input.ndim in (2, 3), "smart_pca: input must be 2D or 3D (D,S or B,D,S)." + assert isinstance(correction, bool), "correction must be bool (True/False)." + + # 2) Convert to batch form if 2D + no_batch = (input.ndim == 2) + if no_batch: + input = input.unsqueeze(0) # shape => (1, D, S) + + B, D, S = input.shape + + # 3) If #dims > #samples, use SVD on (D, S) for each batch + # else use covariance+eigendecomposition + if D > S: + # Possibly center the data + if centered: + mean_ = input.mean(dim=2, keepdim=True) # shape (B, D, 1) + input = input - mean_ + + # We'll collect eigenvalues and eigenvectors for each batch + w_list = [] + v_list = [] + for b in range(B): + inp_ = input[b] # shape (D, S) + # Using SVD => inp_ = U * S_ * V^T + # torch.linalg.svd => (U, S, Vh) + U, Svals, Vh = torch.linalg.svd(inp_, full_matrices=False) + # Convert singular values to eigenvalues + w_ = Svals**2 / (S - (1.0 if correction else 0.0)) + v_ = Vh.transpose(0, 1) # from shape (S, D) => (D, S) + + w_list.append(w_) + v_list.append(v_) + # Stack to shape => (B, D) or (B, S) … but in standard PCA we want dimension = D + w = torch.stack(w_list, dim=0) # shape => (B, k) + v = torch.stack(v_list, dim=0) # shape => (B, D, S) + # If we truly want (B, D, D) we might need to pad or handle the mismatch if D != S + + else: + # Use covariance => (D x D), then eigh + bcov = batch_cov(input, centered=centered, correction=correction) + # bcov shape => (B, D, D) + w_list = [] + v_list = [] + for b in range(B): + C = bcov[b] # shape (D, D) + w_, v_ = eigendecomposition(C, use_rank=use_rank) + w_list.append(w_) + v_list.append(v_) + w = torch.stack(w_list, dim=0) # => shape (B, D) + v = torch.stack(v_list, dim=0) # => shape (B, D, D) + + # 4) If we started with no batch, remove batch dim + if no_batch: + w = w.squeeze(0) # => (D,) + v = v.squeeze(0) # => (D, D) or possibly (D, S) in the SVD branch + + return w, v + +def eigendecomposition(C, use_rank=True): + """ + helper for getting eigenvalues and eigenvectors of covariance matrix + + will measure eigenvalues and eigenvectors with torch.linalg.eigh() + the output will be sorted from highest to lowest eigenvalue (& eigenvector) + + if use_rank=True, will measure the rank of the covariance matrix and zero + out any eigenvalues beyond the rank (that are usually nonzero numerical errors) + """ + try: + w, v = torch.linalg.eigh(C) + except torch._C._LinAlgError as error: + # this happens if the algorithm failed to converge + # try with sklearn's incrementalPCA algorithm + return sklearn_pca(C, use_rank=use_rank) + except Exception as error: + raise error + w_idx = torch.argsort(-w) + w = w[w_idx] + v = v[:, w_idx] + # iff use_rank=True, will set eigenvalues to 0 for probable numerical errors + if use_rank: + crank = torch.linalg.matrix_rank(C) + w[crank:] = 0 + return w, v + +def sklearn_pca(input, use_rank=True, rank=None): + """ + sklearn incrementalPCA algorithm serving as a replacement for eigh when it fails + + input should be a tensor with shape (num_samples, num_features) or it can be a + covariance matrix with (num_features, num_features) + + if use_rank=True, will set num_components to the rank of input and then fill out the + rest of the components with random orthogonal components in the null space of the true + components and set the eigenvalues to 0 + + if use_rank=False, will attempt to fit all the components + if rank is not None, will attempt to fit #=rank components without measuring the rank directly + (will ignore "rank" if use_rank=False) + + returns w, v where w is eigenvalues and v is eigenvectors sorted from highest to lowest + """ + # Move tensor to CPU if it's on GPU + if input.is_cuda: + input = input.cpu() + + num_samples, num_features = input.shape + rank = None if not use_rank else (rank if rank is not None else fast_rank(input)) + ipca = IncrementalPCA(n_components=rank).fit(input) + v = ipca.components_ + w = ipca.singular_values_**2 / num_samples + # if v is a subspace of input (e.g. not a full basis, fill it out) + if v.shape[0] < num_features: + msg = "this condition should always be true, and if not we have to find out why" + assert w.shape[0] == v.shape[0], msg + v_kernel = null_space(v).T + v = np.vstack((v, v_kernel)) + w = np.concatenate((w, np.zeros(v_kernel.shape[0]))) + return torch.tensor(w, dtype=torch.float), torch.tensor(v, dtype=torch.float).T + +def fast_rank(input): + """uses transpose to speed up rank computation, otherwise normal""" + # Move tensor to CPU if it's on GPU + if input.is_cuda: + input = input.cpu() + + if input.size(-2) < input.size(-1): + input = torch.transpose(input, -2, -1) + return int(torch.linalg.matrix_rank(input)) + +def get_maximum_strides(h_input, w_input, layer): + """ + Helper for computing the number of strides h_max, w_max after + convolution with given kernel/stride/padding/dilation. + """ + h_max = int(np.floor((h_input + 2 * layer.padding[0] - layer.dilation[0] * (layer.kernel_size[0] - 1) - 1) / layer.stride[0] + 1)) + w_max = int(np.floor((w_input + 2 * layer.padding[1] - layer.dilation[1] * (layer.kernel_size[1] - 1) - 1) / layer.stride[1] + 1)) + return h_max, w_max + +def get_unfold_params(layer): + return dict(stride=layer.stride, padding=layer.padding, dilation=layer.dilation) + +@torch.no_grad() +def cvPCA(X1, X2): + """X1, X2 are both (dimensions x samples)""" + D, B = X1.shape + assert X2.shape == (D, B), "shape mismatch" + _, u = smart_pca(X1) + cproj0 = X1.T @ u + cproj1 = X2.T @ u + ss = (cproj0 * cproj1).mean(axis=0) + return ss + +def get_num_components(nc, shape): + return nc if nc is not None else min(shape) + +@torch.no_grad() +def shuff_cvPCA(X1, X2, nshuff=5, cvmethod=cvPCA): + D, B = X1.shape + assert X2.shape == (D, B), "shape mismatch" + nc = get_num_components(None, (D, B)) + ss = torch.zeros((nshuff, nc)) + X = torch.stack((X1, X2)) + for k in range(nshuff): + iflip = 1 * (torch.rand(B) > 0.5) + X1c = torch.gather(X, 0, iflip.view(1, 1, -1).expand(1, D, -1)).squeeze(0) + X2c = torch.gather(X, 0, -(iflip - 1).view(1, 1, -1).expand(1, D, -1)).squeeze(0) + ss[k] = cvmethod(X1c, X2c) + return ss + +def avg_value_by_layer(full): + """ + Return average value per layer across a list of epochs or minibatches. + + **full** is a list of lists where the outer list is each snapshot through training or + minibatch etc and each inner list is the value for each node in the network across layers + of a particular measurement + + For example: + num_epochs = 1000 + nodes_per_layer = [50, 40, 30, 20] + len(full) == 1000 + len(full[i]) == 4 ... for all i + [f.shape for f in full[i]] = [50, 40, 30, 20] ... for all i + + this method will return a tensor of size (num_layers, num_epochs) of the average value (for + whatever value is in **full**) for each list/list of values in **full** + """ + num_epochs = len(full) + num_layers = len(full[0]) + avg_full = torch.zeros((num_layers, num_epochs)) + for layer in range(num_layers): + avg_full[layer, :] = torch.tensor([torch.mean(f[layer]) for f in full]) + return avg_full.cpu() + +def value_by_layer(full: List[List[torch.Tensor]], layer: int) -> torch.Tensor: + """ + return all value measurements for a particular layer from **full** + + **full** is a list of lists where the outer list is each snapshot through training or + minibatch etc and each inner list is the value for each node in the network across layers + + this method will return just the part of **full** corresponding to the layer indexed + by **layer** as a tensor of shape (num_epochs, num_nodes) + + see ``avg_value_by_layer`` for a little more explanation + """ + return torch.cat([f[layer].view(1, -1) for f in full], dim=0).cpu() + +def condense_values(full: List[List[List[torch.Tensor]]]) -> List[torch.Tensor]: + """ + condense List[List[List[Tensor]]] -> list of #=num_layers Tensors + shape: (num_networks, num_batches, num_nodes_per_layer) + + returns list of #=num_layers tensors, where each tensor has shape (num_networks, num_batches, num_nodes_per_layer) + + full should be a list of list of lists + the first list should have length = number of networks + the second list should have length = number of batches + the third list should have length = number of layers in the network (this has to be the same for each network!) + the tensor should have shape = number of nodes in this layer (also must be the same for each network) (or can be anything as long as consistent across layers) + """ + num_layers = len(full[0][0]) + return [torch.stack([value_by_layer(value, layer) for value in full]) for layer in range(num_layers)] + +def transpose_list(list_of_lists): + """helper function for transposing the order of a list of lists""" + return list(map(list, zip(*list_of_lists))) + +def named_transpose(list_of_lists, reduction=None): + """ + helper function for transposing lists without forcing the output to be a list like transpose_list + + for example, if list_of_lists contains 10 copies of lists that each have 3 iterable elements you + want to name "A", "B", and "C", then write: + A, B, C = named_transpose(list_of_lists) + + if reduction is used, it will be applied to each output, otherwise will make them lists + """ + if reduction is not None: + return map(reduction, zip(*list_of_lists)) + return map(list, zip(*list_of_lists)) + +def ptp(tensor, dim=None, keepdim=False): + """ + simple method for measuring range of tensor on requested dimension or on all data + """ + if dim is None: + return tensor.max() - tensor.min() + return tensor.max(dim, keepdim).values - tensor.min(dim, keepdim).values + +def rms(tensor, dim=None, keepdim=False): + """simple method for measuring root-mean-square on requested dimension or on all data in tensor""" + if dim is None: + return torch.sqrt(torch.mean(tensor**2)) + return torch.sqrt(torch.mean(tensor**2, dim=dim, keepdim=keepdim)) + +def compute_stats_by_type(tensor, num_types, dim, method="var"): + """ + helper method for returning the mean and variance across a certain dimension + where multiple types are concatenated on that dimension + + for example, suppose we trained 2 networks each with 3 sets of parameters + and concatenated the loss in a tensor like [set1-loss-net1, set1-loss-net2, set2-loss-net1, ...] + then this would contract across the nets from each set and return the mean and variance + """ + num_on_dim = tensor.size(dim) + num_per_type = int(num_on_dim / num_types) + tensor_by_type = tensor.unsqueeze(dim) + expand_shape = list(tensor_by_type.shape) + expand_shape[dim + 1] = num_per_type + expand_shape[dim] = num_types + tensor_by_type = tensor_by_type.view(expand_shape) + type_means = torch.mean(tensor_by_type, dim=dim + 1) + if method == "var": + type_dev = torch.var(tensor_by_type, dim=dim + 1) + elif method == "std": + type_dev = torch.std(tensor_by_type, dim=dim + 1) + elif method == "se": + type_dev = torch.std(tensor_by_type, dim=dim + 1) / np.sqrt(num_per_type) + elif method == "range": + type_dev = ptp(tensor_by_type, dim=dim + 1) + else: + raise ValueError(f"Method ({method}) not recognized.") + return type_means, type_dev + + +def weighted_average(data, weights, dim, keepdim=False, ignore_nan=False): + """ + take the weighted average of **data** on a certain dimension with **weights** + + weights should be a nonnegative vector that broadcasts into data + avg = sum_i(data_i * weight_i, dim) / sum_i(weight_i, dim) + + if ignore_nan=True, (default=False), will ignore nans in weighted average + """ + assert data.ndim == weights.ndim, "data and weights must have same number of dimensions" + assert torch.all(weights[~torch.isnan(weights)] >= 0), "weights must be nonnegative" + + for d in dim if check_iterable(dim) else [dim]: + assert data.size(d) == weights.size( + d + ), f"data and weights must have same size in averaging dimensions (data.size({d})={data.size(d)}, (weight.size({d})={weights.size(d)}))" + + # use normal sum if not ignore nan, otherwise use nansum + sum = torch.nansum if ignore_nan else torch.sum + + # make sure nans are in the same place in weights and data for accurate division by total weight + if ignore_nan: + weights = weights.expand(data.size()) + weights = torch.masked_fill(weights, torch.isnan(data), torch.nan) + + # numerator & denominator of weighted average + numerator = sum(data * weights, dim=dim, keepdim=keepdim) + denominator = sum(weights, dim=dim, keepdim=keepdim) + + # return weighted average + return numerator / denominator + + +def fgsm_attack(image, epsilon, data_grad, transform, sign): + """update an image with fast-gradient sign method""" + warn("fgsm_attack is only going to be in utils temporarily!", DeprecationWarning, stacklevel=2) + # Collect the element-wise sign of the data gradient + if sign: + data_grad = data_grad.sign() + else: + data_grad = data_grad.clone() + # Create the perturbed image by adjusting each pixel of the input image + perturbed_image = image + epsilon * data_grad + # Adding clipping to maintain [0,1] range + perturbed_image = transform(perturbed_image) + # Return the perturbed image + return perturbed_image + +def str2bool(str): + if isinstance(str, bool): + return str + if str.lower() in ("true", "1"): + return True + elif str.lower() in ("false", "0"): + return False + else: + raise TypeError("Boolean type expected") + +def save_checkpoint(nets, optimizers, results, path): + """ + Method for saving checkpoints for networks throughout training. + """ + multi_model_ckpt = {f"model_state_dict_{i}": net.state_dict() for i, net in enumerate(nets)} + multi_optimizer_ckpt = {f"optimizer_state_dict_{i}": opt.state_dict() for i, opt in enumerate(optimizers)} + checkpoint = results | multi_model_ckpt | multi_optimizer_ckpt + torch.save(checkpoint, path) + +def load_checkpoints(nets, optimizers, device, path): + """ + Method for loading presaved checkpoint during training. + """ + if device == "cpu": + checkpoint = torch.load(path, map_location=device) + elif device == "cuda": + checkpoint = torch.load(path) + + net_ids = natsorted([key for key in checkpoint if key.startswith("model_state_dict")]) + opt_ids = natsorted([key for key in checkpoint if key.startswith("optimizer_state_dict")]) + assert all( + [oi.split("_")[-1] == ni.split("_")[-1] for oi, ni in zip(opt_ids, net_ids)] + ), "nets and optimizers cannot be matched up from checkpoint" + + [net.load_state_dict(checkpoint.pop(net_id)) for net, net_id in zip(nets, net_ids)] + [opt.load_state_dict(checkpoint.pop(opt_id)) for opt, opt_id in zip(optimizers, opt_ids)] + + if device == "cuda": + [net.to(device) for net in nets] + return nets, optimizers, checkpoint + +def match_git(path): + """simple method for determining if a path is a git-related file or directory""" + if ".git" in path: + return True + return False + +def compress_directory(output_path, directory_path=None): + """ + Utility to compress entire directory into a zip while respecting .gitignore and ignoring .git. + """ + if directory_path is None: + directory_path = os.path.dirname(os.path.abspath(__file__)) + "/../.." + gitignore_path = os.path.join(directory_path, ".gitignore") + matches = parse_gitignore(gitignore_path) + + files_to_copy = [] + archive_names = [] + for dirpath, dirnames, files in os.walk(directory_path): + if matches(dirpath) or match_git(dirpath): + dirnames[:] = [] + else: + # Filter files based on .gitignore rules (and don't save any .git files) + keep_files = [f for f in files if not matches(f) and not match_git(f)] + full_files = [os.path.join(dirpath, f) for f in keep_files] + for file in full_files: + files_to_copy.append(file) + archive_names.append(os.path.relpath(file, directory_path)) + + with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf: + for file, name in zip(files_to_copy, archive_names): + zipf.write(file, arcname=name) + + +def condense_values(al_list): + data_list = [item["data"] for item in al_list] + aggregated = [] + for net_i in range(len(data_list[0])): + net_snapshots = [d[net_i] for d in data_list] + net_layers = [] + for layer_i in range(len(net_snapshots[0])): + layer_values = [snap[layer_i] for snap in net_snapshots] + net_layers.append(torch.stack(layer_values, dim=0).mean(dim=0)) + aggregated.append(net_layers) + return aggregated \ No newline at end of file diff --git a/_arxiv/_archive/alignment_preref/utils_rem.py b/_arxiv/_archive/alignment_preref/utils_rem.py new file mode 100644 index 00000000..e933d1e2 --- /dev/null +++ b/_arxiv/_archive/alignment_preref/utils_rem.py @@ -0,0 +1,57 @@ +""" +Utilities for alignment research. + +DEPRECATED: This module is kept for backward compatibility. +Please use alignment.utils.core and alignment.utils.math instead. + +This module contains various utility functions for the alignment package. +""" + +import warnings +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +# Import from the new modules for re-export +from alignment.utils.core import ( + setup_logging, + timer, + debug, + to_numpy, + to_tensor, + check_iterable, + ensure_device, + timed +) + +from alignment.utils.math import ( + orthogonalize, + compute_correlation_matrix, + matrix_angles, + project_to_subspace +) + +# Show deprecation warning on import +warnings.warn( + "The alignment.utils module is deprecated and will be removed in a future version. " + "Please use alignment.utils.core and alignment.utils.math instead.", + DeprecationWarning, + stacklevel=2 +) + +# Re-export everything from the imported modules +__all__ = [ + # Core utilities + 'setup_logging', + 'timer', + 'debug', + 'to_numpy', + 'to_tensor', + 'check_iterable', + 'ensure_device', + 'timed', + + # Math utilities + 'orthogonalize', + 'compute_correlation_matrix', + 'matrix_angles', + 'project_to_subspace' +] \ No newline at end of file diff --git a/_arxiv/_archive/alignment_stats.py b/_arxiv/_archive/alignment_stats.py new file mode 100644 index 00000000..1f01906c --- /dev/null +++ b/_arxiv/_archive/alignment_stats.py @@ -0,0 +1,97 @@ +import torch + + +from alignment_v2.datasets import get_dataset +from alignment_v2.models.registry import get_model, get_transform_parameters +from alignment_v2 import processing +from alignment_v2.config import ExperimentConfig +from alignment_v2.experiments.experiment import Experiment +from alignment_v2 import plotting + +class AlignmentStatistics(Experiment): + def get_basename(self): + return "alignment_stats" + + def prepare_path(self): + return [self.args.model.name, self.args.dataset.name, self.args.optimizer.name] + + def create_networks(self): + """ + method for creating networks + """ + if self.args.optimizer.name == "Adam": + optim = torch.optim.Adam + elif self.args.optimizer.name == "SGD": + optim = torch.optim.SGD + else: + raise ValueError(f"optimizer ({self.args.optimizer.name}) not recognized") + + nets = [ + get_model( + self.args.model.name, + alignment_layer_names=self.args.model.alignment_layers, + build=True, + dataset=self.args.dataset.name, + dropout=self.args.model.dropout, + ) + for _ in range(self.args.training.replicates) + ] + nets = [net.to(self.device) for net in nets] + + optimizers = [optim(net.parameters(), lr=self.args.optimizer.lr, weight_decay=self.args.optimizer.weight_decay) for net in nets] + + prms = { + "vals": [self.args.model.name], + "name": "network", + "dataset": self.args.dataset.name, + "dropout": self.args.model.dropout, + "lr": self.args.optimizer.lr, + "weight_decay": self.args.optimizer.weight_decay, + } + return nets, optimizers, prms + + def main(self): + nets, optimizers, prms = self.create_networks() + + dataset = self.prepare_dataset(get_transform_parameters(self.args.model.name, self.args.dataset.name)) + + train_results, test_results = processing.train_networks(self, nets, optimizers, dataset) + + dropout_results, dropout_parameters = processing.progressive_dropout_experiment( + self, nets, dataset, alignment=test_results.get("alignment", None), train_set=False + ) + + eigen_results = processing.measure_eigenfeatures(self, nets, dataset, train_set=False) + + evec_dropout_results, evec_dropout_parameters = processing.eigenvector_dropout(self, nets, dataset, eigen_results, train_set=False) + + results = dict( + prms=prms, + train_results=train_results, + test_results=test_results, + dropout_results=dropout_results, + dropout_parameters=dropout_parameters, + eigen_results=eigen_results, + evec_dropout_results=evec_dropout_results, + evec_dropout_parameters=evec_dropout_parameters, + ) + + return results, nets + + def plot(self, results): + plotting.plot_train_results(self, results["train_results"], results["test_results"], results["prms"]) + plotting.plot_dropout_results( + self, + results["dropout_results"], + results["dropout_parameters"], + results["prms"], + dropout_type="nodes", + ) + plotting.plot_eigenfeatures(self, results["eigen_results"], results["prms"]) + plotting.plot_dropout_results( + self, + results["evec_dropout_results"], + results["evec_dropout_parameters"], + results["prms"], + dropout_type="eigenvectors", + ) \ No newline at end of file diff --git a/_arxiv/_archive/alignment_v2/__init__.py b/_arxiv/_archive/alignment_v2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/_arxiv/_archive/alignment_v2/datasets.py b/_arxiv/_archive/alignment_v2/datasets.py new file mode 100644 index 00000000..5f57be31 --- /dev/null +++ b/_arxiv/_archive/alignment_v2/datasets.py @@ -0,0 +1,322 @@ +from warnings import warn +from abc import ABC, abstractmethod + +import torch +import torchvision +from torch import nn +from torch.utils.data.distributed import DistributedSampler +from torchvision.transforms import v2 as transforms + +from alignment_v2 import files +from alignment_v2.models.base import AlignmentNetwork + +REQUIRED_PROPERTIES = ["dataset_path", "dataset_constructor", "loss_function"] + + +def default_loader_parameters( + distributed, + batch_size=1024*2, + num_workers=2, + shuffle=True, + pin_memory=True, + persistent_workers=True, +): + """ + contains the default dataloader parameters with the option of updating them + using key word argument + """ + default_parameters = dict( + batch_size=batch_size, + num_workers=num_workers, # usually 2 workers is appropriate for swapping loading during batch processing + shuffle=False if distributed else shuffle, # can't use shuffle=True if using DDP + pin_memory=pin_memory, + persistent_workers=persistent_workers, + ) + return default_parameters + + +class DataSet(ABC): + def __init__( + self, + device=None, + distributed=False, + dataset_parameters={}, + transform_parameters={}, + loader_parameters={}, + ): + # set properties of dataset and check that all required properties are defined + self.set_properties() + self.check_properties() + + # define device for dataloading + self.device = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu") + self.distributed = distributed + + # define extra transform (should be a callable method or None) for any transformations that + # can't go in the torchvision.transforms.Compose(...), hopefully this won't be needed later + # when that issue is resolved (grayscale to RGB transform isn't working in Compose right now) + self.extra_transform = transform_parameters.pop("extra_transform", None) + + # create transform for dataloader + self.transform_parameters = transform_parameters + self.make_transform(**transform_parameters) + + # define the dataloader parameters + self.dataloader_parameters = default_loader_parameters(distributed, **loader_parameters) # get dataloader parameters + + # load the dataset and create the dataloaders + self.dataset_parameters = dataset_parameters + self.load_dataset(**dataset_parameters) + + def check_properties(self): + """ + DataSet objects have a few properties that are required but need to be set + by the children. For simplicity, I want the properties to be attributes, + instead of @property methods, but I want to make sure that all the required + properties are loaded. Hence this method. + """ + if not all([hasattr(self, prop) for prop in REQUIRED_PROPERTIES]): + not_found = [prop for prop in REQUIRED_PROPERTIES if not hasattr(self, prop)] + raise ValueError(f"The following required properties were not set: {not_found}") + + @abstractmethod + def set_properties(self): + """ + DataSets have a few required properties (listed in the **REQUIRED_PROPERTIES** + global variable). This method is used to set them all. There is a check implemented + in the __init__ method to make sure all required properties are set. + + required + -------- + dataset_path: string, defines the local file location containing the relevant dataset + files. in particular, whatever filepath is returned by this method will + be passed into the "root" input for torch datasets that is required to + load one of the standard datasets (like MNIST, CIFAR, ImageNet...) + + dataset_constructor: callable method, defines the constructor object for the dataset + loss_function: callable method, defines how to evaluate the loss of the output and target + """ + pass + + @abstractmethod + def dataset_kwargs(self, train=True, **kwargs): + """ + keyword arguments passed into the torch dataset constructor + + different datasets have different kwarg requirements, including + the way train vs test is defined by the kwargs. This class calls + the train dataset and the test dataset using dataset_kwargs(train=True) + or train=False, so the children need to define whatever kwargs go + into the dataset_kwargs appropriately to be determined by the kwarg + train + """ + pass + + def load_dataset(self, **kwargs): + """load dataset using the established path and parameters""" + self.train_dataset = self.dataset_constructor(**self.dataset_kwargs(train=True, **kwargs)) + self.test_dataset = self.dataset_constructor(**self.dataset_kwargs(train=False, **kwargs)) + self.train_sampler = DistributedSampler(self.train_dataset) if self.distributed else None + self.test_sampler = DistributedSampler(self.test_dataset) if self.distributed else None + self.train_loader = torch.utils.data.DataLoader(self.train_dataset, sampler=self.train_sampler, **self.dataloader_parameters) + self.test_loader = torch.utils.data.DataLoader(self.test_dataset, sampler=self.test_sampler, **self.dataloader_parameters) + + def unwrap_batch(self, batch, device=None): + """simple method for unwrapping batch for simple training loops""" + device = self.device if device is None else device + if self.extra_transform: + if type(self.extra_transform) == list: + for et in self.extra_transform: + batch = et(batch) + else: + warn("extra_transform is not a list, this is deprecated!", DeprecationWarning, stacklevel=2) + batch = self.extra_transform(batch) + inputs, targets = batch + inputs, targets = inputs.to(device), targets.to(device) + return inputs, targets + + def make_transform(self, center_crop=None, resize=None, flatten=False, out_channels=None): + """ + create transform for dataloader + resize is the new (H, W) shape of the image for the transforms.Resize transform (or None) + flatten is a boolean indicating whether to flatten the image, (i.e. for a linear input layer) + """ + # default transforms + use_transforms = [ + # Convert PIL Image to PyTorch Tensor + transforms.ToImage(), + transforms.ToDtype(torch.float32, scale=True), + ] + + if center_crop: + use_transforms.append(transforms.CenterCrop(center_crop)) + + # Normalize inputs to canonical distribution + use_transforms.append(transforms.Normalize((self.dist_params["mean"]), (self.dist_params["std"]))) + + # extra transforms depending on network + if resize: + use_transforms.append(transforms.Resize(resize, antialias=True)) + if out_channels: + use_transforms.append(transforms.Grayscale(num_output_channels=out_channels)) + if flatten: + use_transforms.append(transforms.Lambda(torch.flatten)) + + # store composed transformation + self.transform = transforms.Compose(use_transforms) + + def measure_loss(self, outputs, targets, reduction=None): + """simple method for measuring loss with stored loss function""" + if reduction is None: + return self.loss_function(outputs, targets) + + standard_reduction = self.loss_function.reduction + self.loss_function.reduction = reduction + loss = self.loss_function(outputs, targets) + self.loss_function.reduction = standard_reduction + return loss + + def measure_accuracy(self, outputs, targets, k=1, percentage=True): + """ + simple method for measuring accuracy on a classification problem + + default output is top1 percentage, but k can set the top-k accuracy + and if percentage=False then returns the number correct (by topk) + """ + topk = outputs.topk(k, dim=1, sorted=True, largest=True)[1] # get topk indices + num_correct = torch.sum(torch.any(topk == targets.view(-1, 1), dim=1)) # num correct + if percentage: + return 100 * num_correct / outputs.size(0) # percentage + else: + return num_correct + + +class MNIST(DataSet): + def set_properties(self): + """defines the required properties for MNIST""" + self.dataset_path = files.dataset_path("MNIST") + self.dataset_constructor = torchvision.datasets.MNIST + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.1307], std=[0.3081]) + + def dataset_kwargs(self, train=True, download=False): + """set data constructor kwargs for MNIST""" + kwargs = dict( + train=train, + root=self.dataset_path, + download=download, + transform=self.transform, + ) + return kwargs + + +class CIFAR10(DataSet): + def set_properties(self): + """defines the required properties for CIFAR10""" + self.dataset_path = files.dataset_path("CIFAR10") + self.dataset_constructor = torchvision.datasets.CIFAR10 + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + def dataset_kwargs(self, train=True, download=False): + """set data constructor kwargs for CIFAR10""" + kwargs = dict( + train=train, + root=self.dataset_path, + download=download, + transform=self.transform, + ) + return kwargs + + +class CIFAR100(CIFAR10): + def set_properties(self): + """defines the required properties for CIFAR100""" + self.dataset_path = files.dataset_path("CIFAR100") + self.dataset_constructor = torchvision.datasets.CIFAR100 + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + +class ImageNet2012(DataSet): + def set_properties(self): + """ + defines the required properties for ImageNet 2012 (ILSVRC2012) with + 1000 classes. + preprocessing according to pytorch documentation: + https://pytorch.org/hub/pytorch_vision_alexnet/ + """ + self.dataset_path = files.dataset_path("ImageNet") + self.dataset_constructor = torchvision.datasets.ImageNet + self.loss_function = nn.CrossEntropyLoss() + self.dist_params = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) + self.center_crop = 224 + + def dataset_kwargs(self, train=True): + """set data constructor kwargs for ImageNet2012""" + kwargs = dict( + split="train" if train else "val", + root=self.dataset_path, + transform=self.transform, + ) + return kwargs + + +DATASET_REGISTRY = { + "MNIST": MNIST, + "CIFAR10": CIFAR10, + "CIFAR100": CIFAR100, + "ImageNet": ImageNet2012, +} + + +def get_dataset( + dataset_name, + build=False, + dataset_parameters={}, + transform_parameters={}, + loader_parameters={}, + **kwargs, +): + """ + lookup dataset constructor from dataset registry by name + + if build=True, uses kwargs to build dataset and returns a dataset object + otherwise just returns the constructor + """ + if dataset_name not in DATASET_REGISTRY: + raise ValueError(f"Dataset ({dataset_name}) is not in DATASET_REGISTRY") + dataset = DATASET_REGISTRY[dataset_name] + if build: + if isinstance(transform_parameters, AlignmentNetwork): + # Can use an AlignmentNetwork instance to automatically retrieve transform parameters + transform_parameters = transform_parameters.get_transform_parameters(dataset_name) + else: + if not isinstance(transform_parameters, dict): + raise TypeError("transform_parameters must be a dictionary or an AlignmentNetwork") + + # Build the dataset + return dataset( + dataset_parameters=dataset_parameters, + transform_parameters=transform_parameters, + loader_parameters=loader_parameters, + **kwargs, + ) + + # Otherwise return the constructor + return dataset + + +if __name__ == "__main__": + """simple program for downloading a dataset""" + + from argparse import ArgumentParser + + def get_args(args=None): + parser = ArgumentParser(description="simple program for downloading a dataset to the local file location") + parser.add_argument("--dataset", type=str, default="MNIST") + return parser.parse_args(args=args) + + args = get_args() + + dataset = get_dataset(args.dataset, build=True, dataset_parameters=dict(download=True)) diff --git a/_arxiv/_archive/alignment_v2/files.py b/_arxiv/_archive/alignment_v2/files.py new file mode 100644 index 00000000..4aa4216c --- /dev/null +++ b/_arxiv/_archive/alignment_v2/files.py @@ -0,0 +1,55 @@ +import getpass +import socket +from pathlib import Path + +PATH_REGISTRY = { + "DESKTOP-M2J64J2": Path("C:/Users/andrew/Documents/machineLearning"), + "Celia": Path("/Users/celiaberon/Documents/machine_learning"), + "cberon": Path("/n/home00/cberon/alignment/"), + "corwin": Path("/Users/corwin/Building"), + "atlandau": Path("/n/home05/atlandau/machine_learning"), + "landauland": Path("/Users/landauland/Documents/ML-Datasets"), + "nkhoshnevis": Path("/n/vast-scratch/kempner_dev/nkhoshnevis/projects/alignment_dir"), + "hsafaai": Path("/n/holylabs/LABS/kempner_dev/Users/hsafaai/datasets"), +} + + +def get_hostname(): + return socket.gethostname() + + +def get_username(): + return getpass.getuser() + + +def local_path(): + """method for defining the local root path for datasets and results""" + hostname = get_hostname() + if hostname.lower().startswith("celia"): + hostname = "Celia" + hostname = hostname if hostname in PATH_REGISTRY else get_username() + if hostname not in PATH_REGISTRY: + raise ValueError(f"hostname ({hostname}) is not registered in the path registry") + # return path + return PATH_REGISTRY[hostname] + + +def results_path(): + """method for returning relative path to results generated by this package""" + return local_path() / "results" + + +def data_path(): + """method for returning the relative path containing datasets""" + return local_path() / "datasets" + + +def dataset_path(dataset): + """path to specific stored datasets (they all have different requirements)""" + if dataset == "MNIST": + return data_path() + elif dataset == "CIFAR10" or dataset == "CIFAR100": + return data_path() + elif dataset == "ImageNet": + return Path("/n/holylfs06/LABS/kempner_shared/Lab/data/imagenet_1k/") + diff --git a/_arxiv/_archive/alignment_v2/plotting.py b/_arxiv/_archive/alignment_v2/plotting.py new file mode 100644 index 00000000..c9c768e6 --- /dev/null +++ b/_arxiv/_archive/alignment_v2/plotting.py @@ -0,0 +1,565 @@ +import numpy as np +from tqdm import tqdm +from matplotlib import pyplot as plt +import matplotlib as mpl + +import torch +from alignment_v2.utils import compute_stats_by_type, named_transpose, transpose_list, rms + + +def plot_train_results(exp, train_results, test_results, prms): + """ + plotting method for training trajectories and testing data + """ + + num_train_epochs = train_results["loss"].size(0) + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + + print("getting statistics on run data...") + plot_alignment = "alignment" in train_results + if plot_alignment: + alignment = torch.stack([torch.mean(align, dim=2) for align in train_results["alignment"]]) + + cmap = mpl.colormaps["tab10"] + + train_loss_mean, train_loss_se = compute_stats_by_type(train_results["loss"], num_types=num_types, dim=1, method="se") + train_acc_mean, train_acc_se = compute_stats_by_type(train_results["accuracy"], num_types=num_types, dim=1, method="se") + + if plot_alignment: + align_mean, align_se = compute_stats_by_type(alignment, num_types=num_types, dim=1, method="se") + + test_loss_mean, test_loss_se = compute_stats_by_type(torch.tensor(test_results["loss"]), num_types=num_types, dim=0, method="se") + test_acc_mean, test_acc_se = compute_stats_by_type(torch.tensor(test_results["accuracy"]), num_types=num_types, dim=0, method="se") + + print("plotting run data...") + xOffset = [-0.2, 0.2] + get_x = lambda idx: [xOffset[0] + idx, xOffset[1] + idx] + + # Make Training and Testing Performance Figure + alpha = 0.3 + figdim = 3 + figratio = 2 + width_ratios = [figdim, figdim / figratio, figdim, figdim / figratio] + + fig, ax = plt.subplots(1, 4, figsize=(sum(width_ratios), figdim), width_ratios=width_ratios, layout="constrained") + + # plot loss results fot training and testing + for idx, label in enumerate(labels): + cmn = train_loss_mean[:, idx] + cse = train_loss_se[:, idx] + tmn = test_loss_mean[idx] + tse = test_loss_se[idx] + + ax[0].plot(range(num_train_epochs), cmn, color=cmap(idx), label=label) + ax[0].fill_between(range(num_train_epochs), cmn + cse, cmn - cse, color=(cmap(idx), alpha)) + ax[1].plot(get_x(idx), [tmn] * 2, color=cmap(idx), label=label, lw=4) + ax[1].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5) + + ax[0].set_xlabel("Training Epoch") + ax[0].set_ylabel("Loss") + ax[0].set_title("Training Loss") + ax[0].set_ylim(0, None) + ylims = ax[0].get_ylim() + ax[1].set_xticks(range(num_types)) + ax[1].set_xticklabels(labels, rotation=45, ha="right", fontsize=8) + ax[1].set_ylabel("Loss") + ax[1].set_title("Testing") + ax[1].set_xlim(-0.5, num_types - 0.5) + ax[1].set_ylim(ylims) + + # plot loss results fot training and testing + for idx, label in enumerate(labels): + cmn = train_acc_mean[:, idx] + cse = train_acc_se[:, idx] + tmn = test_acc_mean[idx] + tse = test_acc_se[idx] + + ax[2].plot(range(num_train_epochs), cmn, color=cmap(idx), label=label) + ax[2].fill_between(range(num_train_epochs), cmn + cse, cmn - cse, color=(cmap(idx), alpha)) + ax[3].plot(get_x(idx), [tmn] * 2, color=cmap(idx), label=label, lw=4) + ax[3].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5) + + ax[2].set_xlabel("Training Epoch") + ax[2].set_ylabel("Accuracy (%)") + ax[2].set_title("Training Accuracy") + ax[2].set_ylim(0, 100) + ax[3].set_xticks(range(num_types)) + ax[3].set_xticklabels(labels, rotation=45, ha="right", fontsize=8) + ax[3].set_ylabel("Accuracy (%)") + ax[3].set_title("Testing") + ax[3].set_xlim(-0.5, num_types - 0.5) + ax[3].set_ylim(0, 100) + + exp.plot_ready("train_test_performance") + + # Make Alignment Figure + if plot_alignment: + num_align_epochs = align_mean.size(2) + num_layers = align_mean.size(0) + fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", sharex=True) + for idx, label in enumerate(labels): + for layer in range(num_layers): + cmn = align_mean[layer, idx] * 100 + cse = align_se[layer, idx] * 100 + ax[layer].plot(range(num_align_epochs), cmn, color=cmap(idx), label=label) + ax[layer].fill_between(range(num_align_epochs), cmn + cse, cmn - cse, color=(cmap(idx), alpha)) + + for layer in range(num_layers): + ax[layer].set_ylim(0, None) + ax[layer].set_xlabel("Training Epoch") + ax[layer].set_ylabel("Alignment (%)") + ax[layer].set_title(f"Layer {layer}") + + ax[0].legend(loc="lower right") + + exp.plot_ready("train_alignment_by_layer") + + +def plot_dropout_results(exp, dropout_results, dropout_parameters, prms, dropout_type="nodes"): + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val} - dropout {dropout_type}" for val in prms["vals"]] + cmap = mpl.colormaps["Set1"] + alpha = 0.3 + msize = 10 + figdim = 3 + + num_layers = dropout_results["progdrop_loss_high"].size(2) + names = ["From high", "From low", "Random"] + names_diff = ["high-low"] + num_exp = len(names) + dropout_fraction = dropout_results["dropout_fraction"] + by_layer = dropout_results["by_layer"] + extra_name = "by_layer" if by_layer else "all_layers" + extra_name += dropout_type + + # Get statistics across each network type for progressive dropout experiment + print("measuring statistics on dropout analysis...") + loss_mean_high, loss_se_high = compute_stats_by_type(dropout_results["progdrop_loss_high"], num_types=num_types, dim=0, method="se") + loss_mean_low, loss_se_low = compute_stats_by_type(dropout_results["progdrop_loss_low"], num_types=num_types, dim=0, method="se") + loss_mean_rand, loss_se_rand = compute_stats_by_type(dropout_results["progdrop_loss_rand"], num_types=num_types, dim=0, method="se") + + acc_mean_high, acc_se_high = compute_stats_by_type(dropout_results["progdrop_acc_high"], num_types=num_types, dim=0, method="se") + acc_mean_low, acc_se_low = compute_stats_by_type(dropout_results["progdrop_acc_low"], num_types=num_types, dim=0, method="se") + acc_mean_rand, acc_se_rand = compute_stats_by_type(dropout_results["progdrop_acc_rand"], num_types=num_types, dim=0, method="se") + + acc_mean_diff, acc_se_diff = compute_stats_by_type(dropout_results["progdrop_acc_high"]-dropout_results["progdrop_acc_low"], num_types=num_types, dim=0, method="se") + + # Contract into lists for looping through to plot + loss_mean = [loss_mean_high, loss_mean_low, loss_mean_rand] + loss_se = [loss_se_high, loss_se_low, loss_se_rand] + acc_mean = [acc_mean_high, acc_mean_low, acc_mean_rand] + acc_se = [acc_se_high, acc_se_low, acc_se_rand] + + print("plotting dropout results...") + # Plot Loss for progressive dropout experiment + fig, ax = plt.subplots( + num_layers, + num_types, + figsize=(num_types * figdim, num_layers * figdim), + sharex=True, + sharey=True, + layout="constrained", + ) + ax = np.reshape(ax, (num_layers, num_types)) + + if 1==2: + for idx, label in enumerate(labels): + for layer in range(num_layers): + for iexp, name in enumerate(names): + cmn = loss_mean[iexp][idx, :, layer] + cse = loss_se[iexp][idx, :, layer] + ax[layer, idx].plot( + dropout_fraction, + cmn, + color=cmap(iexp), + marker=".", + markersize=msize, + label=name, + ) + ax[layer, idx].fill_between(dropout_fraction, cmn + cse, cmn - cse, color=(cmap(iexp), alpha)) + + if layer == 0: + ax[layer, idx].set_title(label) + + if layer == num_layers - 1: + ax[layer, idx].set_xlabel("Dropout Fraction") + ax[layer, idx].set_xlim(0, 1) + + if idx == 0: + ax[layer, idx].set_ylabel("Loss w/ Dropout") + + if iexp == num_exp - 1: + ax[layer, idx].legend(loc="best") + if exp is not None: + exp.plot_ready("prog_dropout_" + extra_name + "_loss") + + fig, ax = plt.subplots( + num_layers, + num_types, + figsize=(num_types * figdim, num_layers * figdim), + sharex=True, + sharey=True, + layout="constrained", + ) + ax = np.reshape(ax, (num_layers, num_types)) + + for idx, label in enumerate(labels): + for layer in range(num_layers): + for iexp, name in enumerate(names): + + cmn = acc_mean[iexp][idx, :, layer] + cse = acc_se[iexp][idx, :, layer] + ax[layer, idx].plot( + dropout_fraction, + cmn, + color=cmap(iexp), + marker=".", + markersize=msize, + label=name, + ) + ax[layer, idx].fill_between(dropout_fraction, cmn + cse, cmn - cse, color=(cmap(iexp), alpha)) + + ax[layer, idx].set_ylim(0, 100) + + if layer == 0: + ax[layer, idx].set_title(label) + + if layer == num_layers - 1: + ax[layer, idx].set_xlabel("Dropout Fraction") + ax[layer, idx].set_xlim(0, 1) + + if idx == 0: + ax[layer, idx].set_ylabel("Accuracy w/ Dropout") + + if iexp == num_exp - 1: + ax[layer, idx].legend(loc="best") + if exp is not None: + exp.plot_ready("prog_dropout_" + extra_name + "_accuracy") + + # Plot Difference in Accuracy (High - Low) + fig, ax = plt.subplots( + num_layers, + num_types, + figsize=(num_types * figdim, num_layers * figdim), + sharex=True, + sharey=True, + layout="constrained", + ) + ax = np.reshape(ax, (num_layers, num_types)) + + for idx, label in enumerate(labels): + for layer in range(num_layers): + for iexp, name in enumerate(names_diff): + # Debugging: print shapes + #print(f"acc_mean_diff[iexp].shape: {acc_mean_diff[iexp].shape}, dropout_fraction.shape: {dropout_fraction.shape}") + # Squeeze to remove extra dimensions if necessary + cmn = acc_mean_diff[iexp][:, layer] # Squeeze to ensure 1D tensor + cse = acc_se_diff[iexp][:, layer] # Same for standard error + + # Ensure cmn and dropout_fraction have the same length + if len(cmn.shape) == 1 and cmn.shape[0] == dropout_fraction.shape[0]: + ax[layer, idx].plot( + dropout_fraction, + cmn, # Now this should be 1D + color=cmap(iexp), + marker=".", + markersize=msize, + label=name + " (High - Low)", + ) + ax[layer, idx].fill_between(dropout_fraction, cmn + cse, cmn - cse, color=(cmap(iexp), alpha)) + else: + print(f"Shape mismatch: cmn.shape: {cmn.shape}, dropout_fraction.shape: {dropout_fraction.shape}") + + if layer == 0: + ax[layer, idx].set_title(label) + + if layer == num_layers - 1: + ax[layer, idx].set_xlabel("Dropout Fraction") + ax[layer, idx].set_xlim(0, 1) + + if idx == 0: + ax[layer, idx].set_ylabel("Accuracy Difference (High - Low)") + + if iexp == num_exp - 1: + ax[layer, idx].legend(loc="best") + + if exp is not None: + exp.plot_ready("prog_dropout_" + extra_name + "_accuracy_diff") + + +def plot_eigenfeatures(exp, results, prms): + """method for plotting results related to eigen-analysis""" + beta, eigvals, class_betas, class_names = ( + results["beta"], + results["eigvals"], + results["class_betas"], + results["class_names"], + ) + beta = [[torch.abs(b) for b in net_beta] for net_beta in beta] + class_betas = [[rms(cb, dim=2) for cb in net_class_beta] for net_class_beta in class_betas] + + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + cmap = mpl.colormaps["tab10"] + class_cmap = mpl.colormaps["viridis"].resampled(len(class_names)) + + print("measuring statistics of eigenfeature analyses...") + + # shape wrangling + beta = [torch.stack(b) for b in transpose_list(beta)] + eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)] + class_betas = [torch.stack(cb) for cb in transpose_list(class_betas)] + + # normalize to relative values + beta = [b / b.sum(dim=2, keepdim=True) for b in beta] + eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals] + class_betas = [cb / cb.sum(dim=2, keepdim=True) for cb in class_betas] + + # reuse these a few times + statprms = lambda method: dict(num_types=num_types, dim=0, method=method) + + # get mean and variance eigenvalues for each layer for each network type + mean_evals, var_evals = named_transpose([compute_stats_by_type(ev, **statprms("var")) for ev in eigvals]) + + # get sorted betas (sorted within each neuron) + sorted_beta = [torch.sort(b, descending=True, dim=2).values for b in beta] + + # get mean / se beta for each layer for each network type + mean_beta, se_beta = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in beta]) + mean_sorted, se_sorted = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in sorted_beta]) + mean_class_beta, se_class_beta = named_transpose([compute_stats_by_type(cb, **statprms("var")) for cb in class_betas]) + + print("plotting eigenfeature results...") + figdim = 3 + alpha = 0.3 + num_layers = len(mean_beta) + fig, ax = plt.subplots(2, num_layers, figsize=(num_layers * figdim, figdim * 2), layout="constrained") + + for layer in range(num_layers): + num_input = mean_evals[layer].size(1) + num_nodes = mean_beta[layer].size(1) + for idx, label in enumerate(labels): + mn_ev = mean_evals[layer][idx] + se_ev = var_evals[layer][idx] + mn_beta = torch.mean(mean_beta[layer][idx], dim=0) + se_beta = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes) + mn_sort = torch.mean(mean_sorted[layer][idx], dim=0) + se_sort = torch.std(mean_sorted[layer][idx], dim=0) / np.sqrt(num_nodes) + ax[0, layer].plot( + range(num_input), + mn_ev, + color=cmap(idx), + linestyle="--", + label="eigvals" if idx == 0 else None, + ) + ax[0, layer].plot(range(num_input), mn_beta, color=cmap(idx), label=label) + ax[0, layer].fill_between(range(num_input), mn_beta + se_beta, mn_beta - se_beta, color=(cmap(idx), alpha)) + ax[1, layer].plot(range(num_input), mn_sort, color=cmap(idx), label=label) + ax[1, layer].fill_between(range(num_input), mn_sort + se_sort, mn_sort - se_sort, color=(cmap(idx), alpha)) + + ax[0, layer].set_xscale("log") + ax[1, layer].set_xscale("log") + ax[0, layer].set_xlabel("Input Dimension") + ax[1, layer].set_xlabel("Sorted Input Dim") + ax[0, layer].set_ylabel("Relative Eigval & Beta") + ax[1, layer].set_ylabel("Relative Beta (Sorted)") + ax[0, layer].set_title(f"Layer {layer}") + ax[1, layer].set_title(f"Layer {layer}") + + if layer == num_layers - 1: + ax[0, layer].legend(loc="best") + ax[1, layer].legend(loc="best") + + exp.plot_ready("eigenfeatures") + + fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained") + + for layer in range(num_layers): + num_input = mean_evals[layer].size(1) + num_nodes = mean_beta[layer].size(1) + for idx, label in enumerate(labels): + mn_ev = mean_evals[layer][idx] + se_ev = var_evals[layer][idx] + mn_beta = torch.mean(mean_beta[layer][idx], dim=0) + se_beta = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes) + mn_sort = torch.mean(mean_sorted[layer][idx], dim=0) + se_sort = torch.std(mean_sorted[layer][idx], dim=0) / np.sqrt(num_nodes) + ax[layer].plot( + range(num_input), + mn_ev, + color=cmap(idx), + linestyle="--", + label="eigvals" if idx == 0 else None, + ) + ax[layer].plot(range(num_input), mn_beta, color=cmap(idx), label=label) + ax[layer].fill_between(range(num_input), mn_beta + se_beta, mn_beta - se_beta, color=(cmap(idx), alpha)) + + ax[layer].set_xscale("log") + ax[layer].set_yscale("log") + ax[layer].set_xlabel("Input Dimension") + ax[layer].set_ylabel("Relative Eigval & Beta") + ax[layer].set_title(f"Layer {layer}") + + if layer == num_layers - 1: + ax[layer].legend(loc="best") + + exp.plot_ready("eigenfeatures_loglog") + + fig, ax = plt.subplots( + num_types, + num_layers, + figsize=(num_layers * figdim, figdim * num_types), + layout="constrained", + ) + ax = np.reshape(ax, (num_types, num_layers)) + for layer in range(num_layers): + num_input = mean_evals[layer].size(1) + for idx, label in enumerate(labels): + for idx_class, class_name in enumerate(class_names): + mn_data = mean_class_beta[layer][idx][idx_class] + se_data = se_class_beta[layer][idx][idx_class] + ax[idx, layer].plot(range(num_input), mn_data, color=class_cmap(idx_class), label=class_name) + ax[idx, layer].fill_between( + range(num_input), + mn_data + se_data, + mn_data - se_data, + color=(class_cmap(idx_class), alpha), + ) + + ax[idx, layer].set_xscale("log") + ax[idx, layer].set_yscale("linear") + ax[idx, layer].set_xlabel("Input Dimension") + if layer == 0: + ax[idx, layer].set_ylabel(f"{label}\nClass Loading (RMS)") + if idx == 0: + ax[idx, layer].set_title(f"Layer {layer}") + + if layer == num_layers - 1: + ax[idx, layer].legend(loc="upper right", fontsize=6) + + exp.plot_ready("class_eigenfeatures") + + +def plot_adversarial_results(exp, eigen_results, adversarial_results, prms): + accuracy, beta, eigvals = ( + adversarial_results["accuracy"], + adversarial_results["betas"], + eigen_results["eigvals"], + ) + epsilons, use_sign = adversarial_results["epsilons"], adversarial_results["use_sign"] + + num_types = len(prms["vals"]) + labels = [f"{prms['name']}={val}" for val in prms["vals"]] + cmap = mpl.colormaps["tab10"] + + print("measuring statistics of adversarial analyses...") + + # shape wrangling + accuracy = torch.stack([torch.stack(acc) for acc in transpose_list(accuracy)]) # (num_epsilon, num_nets) + eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)] # [(num_nets, dim_layer) for each layer] + beta = [torch.stack(b) for b in beta] # [(epsilon, num_nets, dim_layer) for each layer] + + # normalize to relative values + beta = [b / b.sum(dim=2, keepdim=True) for b in beta] + eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals] + + # reuse these a few times + statprms = lambda dim, method: dict(num_types=num_types, dim=dim, method=method) + + # get mean and variance beta/eigenvalues for each layer for each network type + mean_acc, se_acc = compute_stats_by_type(accuracy, **statprms(1, "var")) # (num_epsilon, num_types) + mean_beta, se_beta = named_transpose( + [compute_stats_by_type(b, **statprms(1, "var")) for b in beta] + ) # [(epsilon, num_types, dim_layer) for each layer] + mean_evals, var_evals = named_transpose( + [compute_stats_by_type(ev, **statprms(0, "var")) for ev in eigvals] + ) # [(num_types, dim_layer) for each layer] + + print("plotting adversarial success results...") + figdim = 3 + alpha = 0.3 + num_layers = len(mean_beta) + fig, ax = plt.subplots(1, 1, figsize=(figdim, figdim), layout="constrained") + for idx, label in enumerate(labels): + ax.plot(epsilons, mean_acc[:, idx], color=cmap(idx), label=label) + ax.set_xlabel("Epsilon") + ax.set_ylabel("Accuracy") + ax.set_title("Adversarial Attack Success") + ax.legend(loc="best") + + exp.plot_ready("adversarial_success") + + print("plotting adversarial structure...") + fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained") + for layer in range(num_layers): + num_input = mean_beta[layer].size(2) + for idx, label in enumerate(labels): + ax[layer].plot( + range(num_input), + torch.nanmean(mean_beta[layer][:, idx], dim=0).detach(), + color=cmap(idx), + label=label, + ) + ax[layer].set_xscale("log") + ax[layer].set_xlabel("Input Dimension") + ax[layer].set_ylabel("Average Component of Pertubation") + ax[layer].set_title(f"Layer {layer}") + if layer == num_layers - 1: + ax[layer].legend(loc="best") + + exp.plot_ready("adversarial_structure") + + +def plot_rf(rf, width, alignment=None, alignBounds=None, showRFs=None, figSize=5): + if showRFs is not None: + rf = rf.reshape(rf.shape[0], -1) + idxRandom = np.random.choice(range(rf.shape[0]), showRFs, replace=False) + rf = rf[idxRandom, :] + else: + showRFs = rf.shape[0] + # normalize + rf = rf.T / np.abs(rf).max(axis=1) + rf = rf.T + rf = rf.reshape(showRFs, width, width) + # If necessary, create colormap + if alignment is not None: + cmap = mpl.cm.get_cmap("rainbow", rf.shape[0]) + cmapPeak = lambda x: cmap(x) + if alignBounds is not None: + alignment = alignment - alignBounds[0] + alignment = alignment / (alignBounds[1] - alignBounds[0]) + else: + alignment = alignment - alignment.min() + alignment = alignment / alignment.max() + + # plotting + n = int(np.ceil(np.sqrt(rf.shape[0]))) + fig, axes = plt.subplots(nrows=n, ncols=n, sharex=True, sharey=True) + fig.set_size_inches(figSize, figSize) + + N = 1000 + for i in tqdm(range(rf.shape[0])): + ax = axes[i // n][i % n] + if alignment is not None: + vals = np.ones((N, 4)) + cAlignment = alignment[i].numpy() + cPeak = cmapPeak(alignment[i].numpy()) + vals[:, 0] = np.linspace(0, cPeak[0], N) + vals[:, 1] = np.linspace(0, cPeak[1], N) + vals[:, 2] = np.linspace(0, cPeak[2], N) + usecmap = mpl.colors.ListedColormap(vals) + ax.imshow(rf[i], cmap=usecmap, vmin=-1, vmax=1) + else: + ax.imshow(rf[i], cmap="gray", vmin=-1, vmax=1) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_aspect("equal") + for j in range(rf.shape[0], n * n): + ax = axes[j // n][j % n] + ax.imshow(np.ones_like(rf[0]) * -1, cmap="gray", vmin=-1, vmax=1) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_aspect("equal") + fig.subplots_adjust(wspace=0.0, hspace=0.0) + return fig diff --git a/_arxiv/_archive/alignment_v2/processing.py b/_arxiv/_archive/alignment_v2/processing.py new file mode 100644 index 00000000..e7db84df --- /dev/null +++ b/_arxiv/_archive/alignment_v2/processing.py @@ -0,0 +1,316 @@ +import os +from tqdm import tqdm +import torch +from alignment_v2 import train +from alignment_v2.utils import load_checkpoints, test_nets, transpose_list, fgsm_attack + + +def train_networks(exp, nets, optimizers, dataset, **special_parameters): + """train and test networks""" + # do training loop + parameters = dict( + train_set=True, + num_epochs=exp.args.epochs, + alignment=not (exp.args.no_alignment), + alignment_expansion=True,#not (exp.args.no_alignment), + delta_weights=exp.args.delta_weights, + frequency=exp.args.frequency, + run=exp.run, + ) + + # update with special parameters + parameters.update(**special_parameters) + + if exp.args.use_prev & os.path.isfile(exp.get_checkpoint_path()): + nets, optimizers, results = load_checkpoints(nets, optimizers, exp.args.device, exp.get_checkpoint_path()) + for net in nets: + net.train() + + parameters["num_complete"] = results["epoch"] + 1 + parameters["results"] = results + print("loaded networks from previous checkpoint") + + if exp.args.save_ckpts: + parameters["save_checkpoints"] = (True, exp.args.ckpt_frequency, exp.get_checkpoint_path(), exp.args.device) + + print("training networks...") + train_results = train.train(nets, optimizers, dataset, **parameters) + + # do testing loop + print("testing networks...") + parameters["train_set"] = False + test_results = train.test(nets, dataset, **parameters) + + return train_results, test_results + + +def progressive_dropout_experiment(exp, nets, dataset, alignment=None, train_set=False): + """ + perform a progressive dropout (of nodes) experiment + alignment is optional, but will be recomputed if you've already measured it. You can provide it + by setting: alignment=test_results['alignment'] if ``train_networks`` has already been run. + """ + # do targeted dropout experiment + print("performing targeted dropout...") + dropout_parameters = dict(num_drops=exp.args.num_drops, by_layer=exp.args.dropout_by_layer, train_set=train_set) + dropout_results = train.progressive_dropout(nets, dataset, alignment=alignment, **dropout_parameters) + return dropout_results, dropout_parameters + +# def progressive_dropout_experiment(exp, nets, dataset, alignment=None, train_set=False, by_layer=False, layer_idx=None): +# """ +# Perform progressive dropout, either layer-by-layer or across all layers at once. +# If by_layer=True, we only dropout nodes in the specified layer (layer_idx). +# """ +# dropout_results = {} +# dropout_parameters = {} + +# if by_layer: +# print(f"Performing dropout for layer {layer_idx}") + +# # Drop nodes in the specified layer only +# dropout_results = train.progressive_dropout( +# nets, dataset, alignment=alignment, layer_idx=layer_idx, train_set=train_set, by_layer=True +# ) + +# else: +# # Standard dropout across all layers +# dropout_results = train.progressive_dropout( +# nets, dataset, alignment=alignment, train_set=train_set, by_layer=False +# ) + +# return dropout_results, dropout_parameters + +def sequential_dropout_experiment(exp, nets, dataset): + """ + Perform a sequential dropout experiment by pruning layers one by one, recalculating alignment after each pruning. + """ + print("Performing sequential dropout experiment...") + sequential_dropout_results = {} + sequential_dropout_parameters = {} + + # Loop over the layers to sequentially prune each layer + for layer_idx in range(len(nets[0].layers)): # Assuming nets[0].layers holds the network layers + print(f"Sequentially pruning layer {layer_idx}") + + # Get the alignment after pruning previous layers + alignment = train.test(nets, dataset, alignment_only=True)["alignment"] + + # Perform progressive dropout for the current layer + print(f"Performing dropout for layer {layer_idx}") + dropout_results, dropout_parameters = progressive_dropout_experiment( + exp, nets, dataset, alignment=alignment, layer_idx=layer_idx + ) + # Recalculate alignment after pruning the current layer + print(f"Recalculating alignment after pruning layer {layer_idx}") + test_results = train.test(nets, dataset, alignment_only=True) # Update test_results after each layer + + # Store the dropout results and parameters for this layer + sequential_dropout_results = dropout_results + sequential_dropout_parameters = dropout_parameters + + return sequential_dropout_results, sequential_dropout_parameters + + +def retrain_network_with_dropout_stats(exp, nets, optimizers, dataset, original_dropout_results): + """ + Retrains the network and maintains the same key structure as the dropout results. + This ensures the retraining results have keys like 'progdrop_loss_high', etc. + """ + # Train the network using the regular training process + print("Retraining the pruned network...") + for net in nets: + net.train() + + # Reset optimizers, assuming we need new ones for the pruned network + optimizers = [torch.optim.SGD(net.parameters(), lr=0.001) for net in nets] + + # Train the networks after dropout + retrain_results, retrain_test_results = train_networks(exp, nets, optimizers, dataset) + + train_results, test_results = train_networks(exp, nets, optimizers, dataset) + + # Initialize retraining results dictionary with same structure as original dropout results + retrain_results = { + "progdrop_loss_high": original_dropout_results["progdrop_loss_high"].clone().zero_(), + "progdrop_loss_low": original_dropout_results["progdrop_loss_low"].clone().zero_(), + "progdrop_loss_rand": original_dropout_results["progdrop_loss_rand"].clone().zero_(), + "progdrop_acc_high": original_dropout_results["progdrop_acc_high"].clone().zero_(), + "progdrop_acc_low": original_dropout_results["progdrop_acc_low"].clone().zero_(), + "progdrop_acc_rand": original_dropout_results["progdrop_acc_rand"].clone().zero_(), + "dropout_fraction": original_dropout_results["dropout_fraction"], + "by_layer": original_dropout_results["by_layer"], + "idx_dropout_layers": original_dropout_results["idx_dropout_layers"], + "fraction_dropped_nodes": original_dropout_results["fraction_dropped_nodes"].clone().zero_() + } + + # Populate the retraining results with new loss and accuracy after retraining + # For example, assuming the train/test functions capture loss and accuracy + retrain_results["progdrop_loss_high"] = test_results.get("progdrop_loss_high", retrain_results["progdrop_loss_high"]) + retrain_results["progdrop_loss_low"] = test_results.get("progdrop_loss_low", retrain_results["progdrop_loss_low"]) + retrain_results["progdrop_loss_rand"] = test_results.get("progdrop_loss_rand", retrain_results["progdrop_loss_rand"]) + + retrain_results["progdrop_acc_high"] = test_results.get("progdrop_acc_high", retrain_results["progdrop_acc_high"]) + retrain_results["progdrop_acc_low"] = test_results.get("progdrop_acc_low", retrain_results["progdrop_acc_low"]) + retrain_results["progdrop_acc_rand"] = test_results.get("progdrop_acc_rand", retrain_results["progdrop_acc_rand"]) + + return retrain_results, test_results + +def evaluate_network(nets, dataset): + """ + Evaluate the pruned network on the test dataset to recalculate alignment and other metrics. + """ + test_results = train.test(nets, dataset) + return test_results + +def measure_eigenfeatures(exp, nets, dataset, train_set=False): + # measure eigenfeatures + print("measuring eigenfeatures...") + beta, eigvals, eigvecs, class_betas = [], [], [], [] + for net in tqdm(nets): + # get inputs to each layer from whole dataloader + real_net = net.module if hasattr(net, "module") else net + inputs, labels = real_net._process_collect_activity( + dataset, + train_set=train_set, + with_updates=False, + use_training_mode=False, + ) + eigenfeatures = real_net.measure_eigenfeatures(inputs, with_updates=False) + beta_by_class = real_net.measure_class_eigenfeatures(inputs, labels, eigenfeatures[2], rms=False, with_updates=False) + beta.append(eigenfeatures[0]) + eigvals.append(eigenfeatures[1]) + eigvecs.append(eigenfeatures[2]) + class_betas.append(beta_by_class) + + + # make it a dictionary + class_names = getattr(dataset.train_loader if train_set else dataset.test_loader, "dataset").classes + return dict( + beta=beta, + eigvals=eigvals, + eigvecs=eigvecs, + class_betas=class_betas, + class_names=class_names, + ) + + +def eigenvector_dropout(exp, nets, dataset, eigen_results, train_set=False): + """ + do targeted eigenvector dropout with precomputed eigenfeatures + """ + # do targeted dropout experiment + print("performing targeted eigenvector dropout...") + evec_dropout_parameters = dict(num_drops=exp.args.num_drops, by_layer=exp.args.dropout_by_layer, train_set=train_set) + evec_dropout_results = train.eigenvector_dropout(nets, dataset, eigen_results["eigvals"], eigen_results["eigvecs"], **evec_dropout_parameters) + return evec_dropout_results, evec_dropout_parameters + + +@test_nets +def measure_adversarial_attacks(nets, dataset, exp, eigen_results, train_set=False, **parameters): + """ + do adversarial attack and measure structure with regards to eigenfeatures + """ + + def get_beta(inputs, eigenvectors): + # get projection of input onto eigenvectors across layers + return [input.cpu() @ evec for input, evec in zip(inputs, eigenvectors)] + + # experiment parameters + epsilons = parameters.get("epsilons") + use_sign = parameters.get("use_sign") + fgsm_transform = parameters.get("fgsm_transform", lambda x: x) + + # data from eigenvectors + eigenvectors = eigen_results["eigvecs"] + + num_eps = len(epsilons) + num_nets = len(nets) + accuracy = torch.zeros((num_nets, num_eps)) + examples = [[[] for _ in range(num_eps)] for _ in range(num_nets)] + betas = [[torch.zeros((num_nets, evec.size(0))) for evec in eigenvectors[0]] for _ in range(num_eps)] + + # dataloader + dataloader = dataset.train_loader if train_set else dataset.test_loader + + for batch in tqdm(dataloader): + input, labels = dataset.unwrap_batch(batch) + + inputs = [input.clone() for _ in range(num_nets)] + + for input in inputs: + input.requires_grad = True + + # Forward pass the data through the model + outputs = [net(input, store_hidden=True) for net, input in zip(nets, inputs)] + input_to_layers = [net.get_layer_inputs(input, precomputed=True) for net in nets] + init_preds = [torch.argmax(output, axis=1) for output in outputs] # find true prediction + least_likely = [torch.argmin(output, axis=1) for output in outputs] # find least likely digit according to model + + c_betas = transpose_list([get_beta(input, evec) for input, evec in zip(input_to_layers, eigenvectors)]) + s_betas = [torch.stack(cb) for cb in c_betas] + + # Calculate the loss + loss = [dataset.measure_loss(output, labels) for output in outputs] + # loss = dataset.measure_loss(output, least_likely) + + # Zero all existing gradients + for net in nets: + net.zero_grad() + + # Calculate gradients of model in backward pass + for l in loss: + l.backward() + + # Collect datagrad + data_grads = [input.grad.data for input in inputs] + + for epsidx, eps in enumerate(epsilons): + + # Call FGSM Attack + perturbed_inputs = [fgsm_attack(input, eps, data_grad, fgsm_transform, use_sign) for input, data_grad in zip(inputs, data_grads)] + + # Re-classify the perturbed image + outputs = [net(perturbed_input, store_hidden=True) for net, perturbed_input in zip(nets, perturbed_inputs)] + input_to_layers = [net.get_layer_inputs(perturbed_input, precomputed=True) for net, perturbed_input in zip(nets, perturbed_inputs)] + c_eps_betas = transpose_list([get_beta(input, evec) for input, evec in zip(input_to_layers, eigenvectors)]) + s_eps_betas = [torch.stack(ceb) for ceb in c_eps_betas] + d_eps_betas = [sebeta - sbeta for sebeta, sbeta in zip(s_eps_betas, s_betas)] + rms_betas = [torch.sqrt(torch.mean(db**2, dim=1)) for db in d_eps_betas] + + for ii, rbeta in enumerate(rms_betas): + betas[epsidx][ii] += rbeta.detach() + + # Check for success + final_preds = [torch.argmax(output, axis=1) for output in outputs] + accuracy[:, epsidx] += torch.tensor([sum(final_pred == labels).cpu() for final_pred in final_preds]) + + # Idx where adversarial example worked + idx_success = [ + torch.where((init_pred == labels) & (final_pred != labels))[0].cpu() for init_pred, final_pred in zip(init_preds, final_preds) + ] + + adv_exs = [perturbed_input.detach().cpu().numpy() for perturbed_input in perturbed_inputs] + for ii, (adv_ex, idx, init_pred, final_pred) in enumerate(zip(adv_exs, idx_success, init_preds, final_preds)): + examples[ii][epsidx].append((init_pred[idx], final_pred[idx], adv_ex[idx])) + + # Calculate final accuracy for this epsilon + accuracy = accuracy / float(len(dataloader.dataset)) + + # Average across betas + betas = transpose_list([[cb / float(len(dataloader.dataset)) for cb in beta] for beta in betas]) + + # Return the accuracy and an adversarial example + return dict(accuracy=accuracy, betas=betas, examples=examples, epsilons=epsilons, use_sign=use_sign) + + +@test_nets +def measure_alignment_distribution(nets, dataset, **parameters): + """ + method for measuring alignment distribution and several associated analyses + """ + # do training loop + parameters = dict( + train_set=True, + ) + + diff --git a/_arxiv/_archive/alignment_v2/train.py b/_arxiv/_archive/alignment_v2/train.py new file mode 100644 index 00000000..0e506501 --- /dev/null +++ b/_arxiv/_archive/alignment_v2/train.py @@ -0,0 +1,813 @@ +from copy import copy, deepcopy + +import torch +from tqdm import tqdm +import torch.distributed as dist # keep for DDP rank checks + +from alignment_v2.utils import ( + transpose_list, + condense_values, + test_nets, + train_nets, + save_checkpoint, + smart_pca, + expected_alignment_distribution, + alignment, + alignment_expansion, +) + + +@train_nets +def train(nets, optimizers, dataset, **parameters): + """method for training network on supervised learning problem""" + + # input argument checks + if not (isinstance(nets, list)): + nets = [nets] + if not (isinstance(optimizers, list)): + optimizers = [optimizers] + assert len(nets) == len(optimizers), "nets and optimizers need to be equal length lists" + + # check if we should print progress bars + verbose = parameters.get("verbose", True) + + num_nets = len(nets) + use_train = parameters.get("train_set", True) + dataloader = dataset.train_loader if use_train else dataset.test_loader + num_steps = len(dataset.train_loader) * parameters["num_epochs"] + + # optional W&B run + run = parameters.get("run") + + # optional analyses + measure_alignment = parameters.get("alignment", True) + measure_alignment_expansion = parameters.get("alignment_expansion", False) + measure_delta_weights = parameters.get("delta_weights", False) + measure_delta_alignment = parameters.get("delta_alignment", False) + measure_frequency = parameters.get("frequency", 1) + compare_expected = parameters.get("compare_expected", False) + + # optional manual shaping + manual_shape = parameters.get("manual_shape", False) + manual_frequency = parameters.get("manual_frequency", -1) + manual_transforms = parameters.get("manual_transforms", None) + manual_layers = parameters.get("manual_layers", None) + + # create or retrieve results dict + results = parameters.get("results", False) + num_complete = parameters.get("num_complete", 0) + save_ckpt, freq_ckpt, path_ckpt, dev = parameters.get("save_checkpoints", (False, 1, "", "")) + if not results: + results = { + "loss": torch.zeros((num_steps, num_nets)), + "accuracy": torch.zeros((num_steps, num_nets)), + } + if measure_alignment: + results["alignment"] = [] + if measure_alignment_expansion: + results["alignment_0"] = [] + results["alignment_1"] = [] + results["alignment_2"] = [] + results["alignment_red"] = [] + if measure_delta_weights: + results["delta_weights"] = [] + results["init_weights"] = [ + net.module.get_alignment_weights() if hasattr(net, "module") else net.get_alignment_weights() + for net in nets + ] + if measure_delta_alignment: + if "init_weights" not in results: + results["init_weights"] = [ + net.module.get_alignment_weights() if hasattr(net, "module") else net.get_alignment_weights() + for net in nets + ] + results["delta_alignment"] = [] + if compare_expected: + calign_bins = torch.linspace(0, 1, 301) + results["compare_alignment_bins"] = calign_bins + results["compare_alignment_expected"] = [] + results["compare_alignment_observed"] = [] + if measure_delta_alignment: + results["compare_delta_alignment_observed"] = [] + elif results["loss"].shape[0] < num_steps: + add_steps = num_steps - results["loss"].shape[0] + assert (add_steps / (parameters["num_epochs"] - num_complete)) == len(dataset.train_loader), \ + "Number of new steps must match epochs * loader length" + results["loss"] = torch.vstack((results["loss"], torch.zeros((add_steps, num_nets)))) + results["accuracy"] = torch.vstack((results["accuracy"], torch.zeros((add_steps, num_nets)))) + + if num_complete > 0: + print("resuming training from checkpoint on epoch", num_complete) + + # training loop + epoch_loop = range(num_complete, parameters["num_epochs"]) + if verbose: + epoch_loop = tqdm(epoch_loop, desc="training epoch") + + for epoch in epoch_loop: + + batch_loop = dataloader + if verbose: + batch_loop = tqdm(batch_loop, desc="minibatch", leave=False) + + for idx, batch in enumerate(batch_loop): + cidx = epoch * len(dataloader) + idx + images, labels = dataset.unwrap_batch(batch) + + for opt in optimizers: + opt.zero_grad() + + outputs = [net(images, store_hidden=True) for net in nets] + loss = [dataset.measure_loss(output, labels) for output in outputs] + for l, opt in zip(loss, optimizers): + l.backward() + opt.step() + + results["loss"][cidx] = torch.tensor([l.item() for l in loss]) + results["accuracy"][cidx] = torch.tensor([ + dataset.measure_accuracy(output, labels).cpu() for output in outputs + ]) + + if idx % measure_frequency == 0: + if measure_alignment: + alignment_vals = [] + for net in nets: + real_net = net.module if hasattr(net, "module") else net + alignment_val = real_net.measure_alignment(images, precomputed=True, method="alignment") + alignment_vals.append(alignment_val) + results["alignment"].append(alignment_vals) + + if measure_alignment_expansion: + alignment0_vals, alignment1_vals, alignment2_vals, alignmentred_vals = [], [], [], [] + for net in nets: + real_net = net.module if hasattr(net, "module") else net + alignment0_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_0")) + alignment1_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_1")) + alignment2_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_2")) + alignmentred_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_red")) + results["alignment_0"].append(alignment0_vals) + results["alignment_1"].append(alignment1_vals) + results["alignment_2"].append(alignment2_vals) + results["alignment_red"].append(alignmentred_vals) + + if measure_delta_weights or measure_delta_alignment: + c_delta_weights = [] + for net, init_weight in zip(nets, results["init_weights"]): + real_net = net.module if hasattr(net, "module") else net + c_delta_weights.append(real_net.compare_weights(init_weight)) + if measure_delta_weights: + results["delta_weights"].append(c_delta_weights) + if measure_delta_alignment: + c_delta_alignment = [] + for net, weights_ in zip(nets, c_delta_weights): + real_net = net.module if hasattr(net, "module") else net + c_delta_alignment.append( + real_net.measure_alignment_weights(images, weights_, precomputed=True, method="alignment") + ) + results["delta_alignment"].append(c_delta_alignment) + + if compare_expected: + if measure_alignment: + c_alignment = results["alignment"][-1] + else: + c_alignment = [] + for net in nets: + real_net = net.module if hasattr(net, "module") else net + c_alignment.append(real_net.measure_alignment(images, precomputed=True, method="alignment")) + + c_inputs = [] + for net in nets: + real_net = net.module if hasattr(net, "module") else net + c_inputs.append(real_net.get_layer_inputs(images, precomputed=True)) + + for i in range(len(nets)): + real_net_ = nets[i].module if hasattr(nets[i], "module") else nets[i] + c_inputs[i] = real_net_._preprocess_inputs(c_inputs[i]) + + c_evals = [] + for cin in c_inputs: + c_eval_sub = [] + for c_ in cin: + w, _ = smart_pca(c_.T) + c_eval_sub.append(w) + c_evals.append(c_eval_sub) + + calign_bins = results["compare_alignment_bins"] + c_dist = [] + for c_eval in c_evals: + subdist = [] + for ev in c_eval: + subdist.append(expected_alignment_distribution(ev, valid_rotation=False, bins=calign_bins)[0]) + c_dist.append(subdist) + + t_dist = [] + for c_align in c_alignment: + align_sub = [] + for align in c_align: + align_sub.append(torch.histogram(align.cpu(), bins=calign_bins, density=True)[0]) + t_dist.append(align_sub) + + results["compare_alignment_expected"].append(c_dist) + results["compare_alignment_observed"].append(t_dist) + + if measure_delta_alignment: + d_alignment = results["delta_alignment"][-1] + d_dist = [] + for d_align in d_alignment: + align_sub = [] + for dalign in d_align: + align_sub.append(torch.histogram(dalign.cpu(), bins=calign_bins, density=True)[0]) + d_dist.append(align_sub) + results["compare_delta_alignment_observed"].append(d_dist) + + # Only rank 0 logs to W&B + if run is not None: + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + pass + else: + run.log( + {f"losses/loss-{ii}": l.item() for ii, l in enumerate(loss)} + | {f"accuracies/accuracy-{ii}": dataset.measure_accuracy(output, labels) for ii, output in enumerate(outputs)} + | {"batch": cidx} + ) + + if manual_shape: + if ((epoch + 1) % manual_frequency == 0) and (epoch < parameters["num_epochs"] - 1): + for net, transform in tqdm(zip(nets, manual_transforms), desc="manual shaping", leave=False): + real_net = net.module if hasattr(net, "module") else net + inputs, _ = real_net._process_collect_activity(dataset, train_set=False, with_updates=False, use_training_mode=False) + _, eigenvalues, eigenvectors = real_net.measure_eigenfeatures(inputs, with_updates=False) + idx_to_layer_lookup = {layer: i for i, layer in enumerate(real_net.get_alignment_layer_indices())} + eigenvalues = [eigenvalues[idx_to_layer_lookup[ml]] for ml in manual_layers] + eigenvectors = [eigenvectors[idx_to_layer_lookup[ml]] for ml in manual_layers] + real_net.shape_eigenfeatures(manual_layers, eigenvalues, eigenvectors, transform) + + if save_ckpt & (epoch % freq_ckpt == 0): + save_checkpoint( + nets, + optimizers, + results | {"prms": parameters, "epoch": epoch, "device": dev}, + path_ckpt, + ) + + for k in [ + "alignment", + "alignment_0", + "alignment_1", + "alignment_2", + "alignment_red", + "delta_weights", + "delta_alignment", + "avgcorr", + "fullcorr", + "compare_alignment_expected", + "compare_alignment_observed", + "compare_delta_alignment_observed", + ]: + if k not in results.keys(): + continue + results[k] = condense_values(transpose_list(results[k])) + + return results + + + +@torch.no_grad() +@test_nets +def test(nets, dataset, **parameters): + """method for testing network on supervised learning problem""" + + run = parameters.get("run") + + if not (isinstance(nets, list)): + nets = [nets] + + verbose = parameters.get("verbose", True) + num_nets = len(nets) + + use_test = not parameters.get("train_set", False) + dataloader = dataset.test_loader if use_test else dataset.train_loader + + total_loss = [0 for _ in range(num_nets)] + num_correct = [0 for _ in range(num_nets)] + num_batches = 0 + + measure_alignment = parameters.get("alignment", True) + measure_alignment_expansion = parameters.get("alignment_expansion", True) + + if measure_alignment: + alignment = [] + if measure_alignment_expansion: + alignment_0 = [] + alignment_1 = [] + alignment_2 = [] + alignment_red = [] + + batch_loop = tqdm(dataloader) if verbose else dataloader + for batch in batch_loop: + images, labels = dataset.unwrap_batch(batch) + + outputs = [net(images, store_hidden=True) for net in nets] + + for idx, output in enumerate(outputs): + total_loss[idx] += dataset.measure_loss(output, labels).item() + num_correct[idx] += dataset.measure_accuracy(output, labels).item() + + num_batches += 1 + + if measure_alignment: + a_vals = [] + for net in nets: + real_net = net.module if hasattr(net, "module") else net + a_vals.append(real_net.measure_alignment(images, precomputed=True, method="alignment")) + alignment.append(a_vals) + + if measure_alignment_expansion: + a0_vals, a1_vals, a2_vals, ared_vals = [], [], [], [] + for net in nets: + real_net = net.module if hasattr(net, "module") else net + a0_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_0")) + a1_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_1")) + a2_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_2")) + ared_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_red")) + alignment_0.append(a0_vals) + alignment_1.append(a1_vals) + alignment_2.append(a2_vals) + alignment_red.append(ared_vals) + + results = { + "loss": [loss / num_batches for loss in total_loss], + "accuracy": [correct / num_batches for correct in num_correct], + } + + if measure_alignment: + results["alignment"] = condense_values(transpose_list(alignment)) + if measure_alignment_expansion: + results["alignment_0"] = condense_values(transpose_list(alignment_0)) + results["alignment_1"] = condense_values(transpose_list(alignment_1)) + results["alignment_2"] = condense_values(transpose_list(alignment_2)) + results["alignment_red"] = condense_values(transpose_list(alignment_red)) + + if run is not None: + run.summary["test_loss"] = torch.mean(torch.tensor(results["loss"])) + run.summary["test_accuracy"] = torch.mean(torch.tensor(results["accuracy"])) + + return results + + +@torch.no_grad() +def get_dropout_indices(idx_alignment, fraction): + """ + convenience method for getting a fraction of dropout indices from each layer + """ + num_nets = idx_alignment[0].size(0) + num_nodes = [idx.size(1) for idx in idx_alignment] + num_drop = [int(nodes * fraction) for nodes in num_nodes] + idx_high = [ + torch.sort(idx[:, -drop:], dim=1).values + for idx, drop in zip(idx_alignment, num_drop) + ] + idx_low = [ + torch.sort(idx[:, :drop], dim=1).values + for idx, drop in zip(idx_alignment, num_drop) + ] + idx_rand = [ + torch.sort(idx[:, torch.randperm(idx.size(1))[:drop]], dim=1).values + for idx, drop in zip(idx_alignment, num_drop) + ] + return idx_high, idx_low, idx_rand + + +@torch.no_grad() +def get_dropout_indices000(idx_alignment, fraction): + """ + Returns dropout indices for high, low, and random alignment for each layer and each network. + """ + num_layers = len(idx_alignment) + num_nets = idx_alignment[0].size(0) + + idx_high = [] + idx_low = [] + idx_rand = [] + + for layer_idx in range(num_layers): + layer_alignment = idx_alignment[layer_idx] + layer_dimension = layer_alignment.size(1) + num_drop = int(layer_dimension * fraction) + sorted_indices = torch.argsort(layer_alignment, dim=1) + high_indices = sorted_indices[:, -num_drop:] + high_indices = torch.sort(high_indices, dim=1).values + idx_high.append(high_indices) + low_indices = sorted_indices[:, :num_drop] + low_indices = torch.sort(low_indices, dim=1).values + idx_low.append(low_indices) + rand_indices = torch.randperm(layer_dimension, device=layer_alignment.device)[:num_drop] + rand_indices = rand_indices.repeat(num_nets, 1) + idx_rand.append(rand_indices) + + return idx_high, idx_low, idx_rand + + +@torch.no_grad() +@test_nets +def progressive_dropout(nets, dataset, alignment=None, **parameters): + """ + method for testing network on supervised learning problem with progressive dropout + """ + + if not (isinstance(nets, list)): + nets = [nets] + + real_net = nets[0].module if hasattr(nets[0], 'module') else nets[0] + idx_dropout_layers = real_net.get_alignment_layer_indices() + + if alignment is None: + alignment = test(nets, dataset, **parameters)["alignment"] + + alignment = [alignment[i] for i in idx_dropout_layers] + assert len(alignment) == len(idx_dropout_layers), "the number of layers in **alignment** doesn't correspond to the number of alignment layers" + + classification_layer = ( + nets[0].module.num_layers(all=True) - 1 + if hasattr(nets[0], 'module') + else nets[0].num_layers(all=True) - 1 + ) + if classification_layer in idx_dropout_layers: + idx_dropout_layers.pop(-1) + alignment.pop(-1) + + alignment = [torch.mean(align, dim=1) for align in alignment] + idx_alignment = [torch.argsort(align, dim=1) for align in alignment] + + print(len(idx_alignment)) + print(idx_alignment[0].shape) + + num_nets = len(nets) + num_drops = parameters.get("num_drops", 9) + drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1] + by_layer = parameters.get("by_layer", False) + num_layers = len(idx_dropout_layers) if by_layer else 1 + + progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers)) + + num_batches = 0 + use_train = parameters.get("train_set", False) + dataloader = dataset.train_loader if use_train else dataset.test_loader + + for batch in tqdm(dataloader): + images, labels = dataset.unwrap_batch(batch) + num_batches += 1 + + for dropidx, fraction in enumerate(drop_fraction): + idx_high, idx_low, idx_rand = get_dropout_indices(idx_alignment, fraction) + + for layer in range(num_layers): + if by_layer: + drop_high, drop_low, drop_rand = ( + [idx_high[layer]], + [idx_low[layer]], + [idx_rand[layer]], + ) + drop_layer = [idx_dropout_layers[layer]] + else: + drop_high, drop_low, drop_rand = idx_high, idx_low, idx_rand + drop_layer = copy(idx_dropout_layers) + + out_high = [] + for idx, net_ in enumerate(nets): + real_net_ = net_.module if hasattr(net_, "module") else net_ + out_ = real_net_.forward_targeted_dropout( + images, [drop[idx, :] for drop in drop_high], drop_layer + )[0] + out_high.append(out_) + out_low = [] + for idx, net_ in enumerate(nets): + real_net_ = net_.module if hasattr(net_, "module") else net_ + out_ = real_net_.forward_targeted_dropout( + images, [drop[idx, :] for drop in drop_low], drop_layer + )[0] + out_low.append(out_) + out_rand = [] + for idx, net_ in enumerate(nets): + real_net_ = net_.module if hasattr(net_, "module") else net_ + out_ = real_net_.forward_targeted_dropout( + images, [drop[idx, :] for drop in drop_rand], drop_layer + )[0] + out_rand.append(out_) + + loss_high = [dataset.measure_loss(out, labels).item() for out in out_high] + loss_low = [dataset.measure_loss(out, labels).item() for out in out_low] + loss_rand = [dataset.measure_loss(out, labels).item() for out in out_rand] + + acc_high = [dataset.measure_accuracy(out, labels) for out in out_high] + acc_low = [dataset.measure_accuracy(out, labels) for out in out_low] + acc_rand = [dataset.measure_accuracy(out, labels) for out in out_rand] + + progdrop_loss_high[:, dropidx, layer] += torch.tensor(loss_high) + progdrop_loss_low[:, dropidx, layer] += torch.tensor(loss_low) + progdrop_loss_rand[:, dropidx, layer] += torch.tensor(loss_rand) + progdrop_acc_high[:, dropidx, layer] += torch.tensor(acc_high) + progdrop_acc_low[:, dropidx, layer] += torch.tensor(acc_low) + progdrop_acc_rand[:, dropidx, layer] += torch.tensor(acc_rand) + + results = { + "progdrop_loss_high": progdrop_loss_high / num_batches, + "progdrop_loss_low": progdrop_loss_low / num_batches, + "progdrop_loss_rand": progdrop_loss_rand / num_batches, + "progdrop_acc_high": progdrop_acc_high / num_batches, + "progdrop_acc_low": progdrop_acc_low / num_batches, + "progdrop_acc_rand": progdrop_acc_rand / num_batches, + "dropout_fraction": drop_fraction, + "by_layer": by_layer, + "idx_dropout_layers": idx_dropout_layers, + } + + return results + + +def train_network(net, dataset, epochs=5, learning_rate=0.001): + optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate) + net.train() + + for epoch in range(epochs): + for batch in dataset.train_loader: + images, labels = dataset.unwrap_batch(batch) + optimizer.zero_grad() + outputs = net(images) + loss = dataset.measure_loss(outputs, labels) + loss.backward() + optimizer.step() + + return net + + +def progressive_dropout_train(nets, dataset, alignment=None, **parameters): + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if not isinstance(nets, list): + nets = [nets] + nets = [net.to(device) for net in nets] + + real_net_0 = nets[0].module if hasattr(nets[0], 'module') else nets[0] + idx_dropout_layers = real_net_0.get_alignment_layer_indices() + + if alignment is None: + alignment = test(nets, dataset, **parameters)["alignment"] + + assert len(alignment) == len(idx_dropout_layers), ( + f"Mismatch in alignment layer count: alignment has {len(alignment)} " + f"layers, expected {len(idx_dropout_layers)} layers." + ) + + classification_layer = real_net_0.num_layers(all=True) - 1 + if classification_layer in idx_dropout_layers: + idx_dropout_layers.pop(-1) + alignment.pop(-1) + + alignment = [torch.mean(align, dim=1) for align in alignment] + idx_alignment = [torch.argsort(align, dim=1) for align in alignment] + + num_nets = len(nets) + num_drops = parameters.get("num_drops", 9) + drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1] + by_layer = parameters.get("by_layer", False) + num_layers = len(idx_dropout_layers) if by_layer else 1 + + progdrop_loss_high_before, progdrop_loss_low_before, progdrop_loss_rand_before = [], [], [] + progdrop_acc_high_before, progdrop_acc_low_before, progdrop_acc_rand_before = [], [], [] + + progdrop_loss_high_after, progdrop_loss_low_after, progdrop_loss_rand_after = [], [], [] + progdrop_acc_high_after, progdrop_acc_low_after, progdrop_acc_rand_after = [], [], [] + + num_batches = 0 + use_train = parameters.get("train_set", False) + dataloader = dataset.train_loader if use_train else dataset.test_loader + + for batch in tqdm(dataloader): + images, labels = dataset.unwrap_batch(batch) + images = images.to(device) + labels = labels.to(device) + + num_batches += 1 + + for dropidx, fraction in enumerate(drop_fraction): + idx_high, idx_low, idx_rand = get_dropout_indices(idx_alignment, fraction) + + for layer in range(num_layers): + if by_layer: + drop_high, drop_low, drop_rand = ( + [idx_high[layer]], + [idx_low[layer]], + [idx_rand[layer]], + ) + drop_layer = [idx_dropout_layers[layer]] + else: + drop_high, drop_low, drop_rand = idx_high, idx_low, idx_rand + drop_layer = deepcopy(idx_dropout_layers) + + pruned_nets_high = [ + train_network(prune_network_by_alignment([net], idx_alignment, fraction, method="high")[0].to(device), dataset) + for net in nets + ] + pruned_nets_low = [ + train_network(prune_network_by_alignment([net], idx_alignment, fraction, method="low")[0].to(device), dataset) + for net in nets + ] + pruned_nets_rand = [ + train_network(prune_network_by_alignment([net], idx_alignment, fraction, method="random")[0].to(device), dataset) + for net in nets + ] + + out_high_before = [net_(images) for net_ in pruned_nets_high] + out_low_before = [net_(images) for net_ in pruned_nets_low] + out_rand_before = [net_(images) for net_ in pruned_nets_rand] + + loss_high_before = [ + dataset.measure_loss(out, labels).item() for out in out_high_before + ] + loss_low_before = [ + dataset.measure_loss(out, labels).item() for out in out_low_before + ] + loss_rand_before = [ + dataset.measure_loss(out, labels).item() for out in out_rand_before + ] + + acc_high_before = [dataset.measure_accuracy(out, labels) for out in out_high_before] + acc_low_before = [dataset.measure_accuracy(out, labels) for out in out_low_before] + acc_rand_before = [dataset.measure_accuracy(out, labels) for out in out_rand_before] + + progdrop_loss_high_before.append(loss_high_before) + progdrop_loss_low_before.append(loss_low_before) + progdrop_loss_rand_before.append(loss_rand_before) + progdrop_acc_high_before.append(acc_high_before) + progdrop_acc_low_before.append(acc_low_before) + progdrop_acc_rand_before.append(acc_rand_before) + + out_high_after = [net_(images) for net_ in pruned_nets_high] + out_low_after = [net_(images) for net_ in pruned_nets_low] + out_rand_after = [net_(images) for net_ in pruned_nets_rand] + + loss_high_after = [ + dataset.measure_loss(out, labels).item() for out in out_high_after + ] + loss_low_after = [ + dataset.measure_loss(out, labels).item() for out in out_low_after + ] + loss_rand_after = [ + dataset.measure_loss(out, labels).item() for out in out_rand_after + ] + + acc_high_after = [dataset.measure_accuracy(out, labels) for out in out_high_after] + acc_low_after = [dataset.measure_accuracy(out, labels) for out in out_low_after] + acc_rand_after = [dataset.measure_accuracy(out, labels) for out in out_rand_after] + + progdrop_loss_high_after.append(loss_high_after) + progdrop_loss_low_after.append(loss_low_after) + progdrop_loss_rand_after.append(loss_rand_after) + progdrop_acc_high_after.append(acc_high_after) + progdrop_acc_low_after.append(acc_low_after) + progdrop_acc_rand_after.append(acc_rand_after) + + results = { + "progdrop_loss_high_before": torch.tensor(progdrop_loss_high_before) / num_batches, + "progdrop_loss_low_before": torch.tensor(progdrop_loss_low_before) / num_batches, + "progdrop_loss_rand_before": torch.tensor(progdrop_loss_rand_before) / num_batches, + "progdrop_acc_high_before": torch.tensor(progdrop_acc_high_before) / num_batches, + "progdrop_acc_low_before": torch.tensor(progdrop_acc_low_before) / num_batches, + "progdrop_acc_rand_before": torch.tensor(progdrop_acc_rand_before) / num_batches, + "progdrop_loss_high_after": torch.tensor(progdrop_loss_high_after) / num_batches, + "progdrop_loss_low_after": torch.tensor(progdrop_loss_low_after) / num_batches, + "progdrop_loss_rand_after": torch.tensor(progdrop_loss_rand_after) / num_batches, + "progdrop_acc_high_after": torch.tensor(progdrop_acc_high_after) / num_batches, + "progdrop_acc_low_after": torch.tensor(progdrop_acc_low_after) / num_batches, + "progdrop_acc_rand_after": torch.tensor(progdrop_acc_rand_after) / num_batches, + "dropout_fraction": drop_fraction, + "by_layer": by_layer, + "idx_dropout_layers": idx_dropout_layers, + } + + return results + + +@torch.no_grad() +@test_nets +def eigenvector_dropout(nets, dataset, eigenvalues, eigenvectors, **parameters): + """ + method for testing network on supervised learning problem with eigenvector dropout + """ + + if not (isinstance(nets, list)): + nets = [nets] + + idx_dropout_layers = ( + nets[0].module.get_alignment_layer_indices() + if hasattr(nets[0], 'module') + else nets[0].get_alignment_layer_indices() + ) + + assert all([len(ev) == len(idx_dropout_layers) for ev in eigenvectors]), \ + "the number of layers in **eigenvectors** doesn't correspond to the number of alignment layers" + assert all([len(ev) == len(idx_dropout_layers) for ev in eigenvalues]), \ + "the number of layers in **eigenvalues** doesn't correspond to the number of alignment layers" + + num_nets = len(nets) + num_drops = parameters.get("num_drops", 9) + drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1] + by_layer = parameters.get("by_layer", False) + num_layers = len(idx_dropout_layers) if by_layer else 1 + + idx_eigenvalue = [ + torch.fliplr(torch.tensor(range(0, ev.size(1))).expand(num_nets, -1)) + for ev in eigenvectors[0] + ] + + progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers)) + progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers)) + + num_batches = 0 + use_test = not parameters.get("train_set", True) + dataloader = dataset.test_loader if use_test else dataset.train_loader + + for batch in tqdm(dataloader): + images, labels = dataset.unwrap_batch(batch) + num_batches += 1 + + for dropidx, fraction in enumerate(drop_fraction): + idx_high, idx_low, idx_rand = get_dropout_indices(idx_eigenvalue, fraction) + + for layer in range(num_layers): + if by_layer: + drop_high, drop_low, drop_rand = ( + [idx_high[layer]], + [idx_low[layer]], + [idx_rand[layer]], + ) + drop_layer = [idx_dropout_layers[layer]] + drop_evals = [[ev[layer]] for ev in eigenvalues] + drop_evecs = [[evc[layer]] for evc in eigenvectors] + else: + drop_high, drop_low, drop_rand = idx_high, idx_low, idx_rand + drop_layer = copy(idx_dropout_layers) + drop_evals = deepcopy(eigenvalues) + drop_evecs = deepcopy(eigenvectors) + + out_high = [] + for idx, (net, evals, evecs) in enumerate(zip(nets, drop_evals, drop_evecs)): + real_net = net.module if hasattr(net, 'module') else net + out_ = real_net.forward_eigenvector_dropout( + images, evals, evecs, [drop[idx, :] for drop in drop_high], drop_layer + )[0] + out_high.append(out_) + out_low = [] + for idx, (net, evals, evecs) in enumerate(zip(nets, drop_evals, drop_evecs)): + real_net = net.module if hasattr(net, 'module') else net + out_ = real_net.forward_eigenvector_dropout( + images, evals, evecs, [drop[idx, :] for drop in drop_low], drop_layer + )[0] + out_low.append(out_) + out_rand = [] + for idx, (net, evals, evecs) in enumerate(zip(nets, drop_evals, drop_evecs)): + real_net = net.module if hasattr(net, 'module') else net + out_ = real_net.forward_eigenvector_dropout( + images, evals, evecs, [drop[idx, :] for drop in drop_rand], drop_layer + )[0] + out_rand.append(out_) + + loss_high = [dataset.measure_loss(out, labels).item() for out in out_high] + loss_low = [dataset.measure_loss(out, labels).item() for out in out_low] + loss_rand = [dataset.measure_loss(out, labels).item() for out in out_rand] + + acc_high = [dataset.measure_accuracy(out, labels) for out in out_high] + acc_low = [dataset.measure_accuracy(out, labels) for out in out_low] + acc_rand = [dataset.measure_accuracy(out, labels) for out in out_rand] + + progdrop_loss_high[:, dropidx, layer] += torch.tensor(loss_high) + progdrop_loss_low[:, dropidx, layer] += torch.tensor(loss_low) + progdrop_loss_rand[:, dropidx, layer] += torch.tensor(loss_rand) + progdrop_acc_high[:, dropidx, layer] += torch.tensor(acc_high) + progdrop_acc_low[:, dropidx, layer] += torch.tensor(acc_low) + progdrop_acc_rand[:, dropidx, layer] += torch.tensor(acc_rand) + + results = { + "progdrop_loss_high": progdrop_loss_high / num_batches, + "progdrop_loss_low": progdrop_loss_low / num_batches, + "progdrop_loss_rand": progdrop_loss_rand / num_batches, + "progdrop_acc_high": progdrop_acc_high / num_batches, + "progdrop_acc_low": progdrop_acc_low / num_batches, + "progdrop_acc_rand": progdrop_acc_rand / num_batches, + "dropout_fraction": drop_fraction, + "by_layer": by_layer, + "idx_dropout_layers": idx_dropout_layers, + } + + return results \ No newline at end of file diff --git a/_arxiv/_archive/alignment_v2/utils.py b/_arxiv/_archive/alignment_v2/utils.py new file mode 100644 index 00000000..73149842 --- /dev/null +++ b/_arxiv/_archive/alignment_v2/utils.py @@ -0,0 +1,934 @@ +import os +import math +import zipfile +from typing import List +from warnings import warn +from contextlib import contextmanager +from functools import wraps +from natsort import natsorted +from gitignore_parser import parse_gitignore + +import torch +import numpy as np +from scipy.linalg import null_space +from sklearn.decomposition import IncrementalPCA +from sklearn.decomposition import PCA + +# -------------- context managers & decorators -------------- +@contextmanager +def no_grad(no_grad=True): + if no_grad: + with torch.no_grad(): + yield + else: + yield + + +def test_nets(func): + @wraps(func) + def wrapper(nets, *args, **kwargs): + # get original training mode and set to eval + in_training_mode = [set_net_mode(net, training=False) for net in nets] + + # do decorated function + func_outputs = func(nets, *args, **kwargs) + + # return networks to whatever mode they used to be in + for train_mode, net in zip(in_training_mode, nets): + set_net_mode(net, training=train_mode) + + # return decorated function outputs + return func_outputs + + # return decorated function + return wrapper + + +def train_nets(func): + @wraps(func) + def wrapper(nets, *args, **kwargs): + # get original training mode and set to train + in_training_mode = [set_net_mode(net, training=True) for net in nets] + + # do decorated function + func_outputs = func(nets, *args, **kwargs) + + # return networks to whatever mode they used to be in + for train_mode, net in zip(in_training_mode, nets): + set_net_mode(net, training=train_mode) + + # return decorated function outputs + return func_outputs + + # return decorated function + return wrapper + + +def set_net_mode(net, training=True): + """helper for setting mode of network and returning current mode""" + # get current mode of network + in_training_mode = net.training + # set to training mode or evaluation mode + if training: + net.train() + else: + net.eval() + # return original mode of network + return in_training_mode + + +# ------- some other things ------- +def get_device(obj): + """simple method to get device of input tensor or nn.Module""" + if isinstance(obj, torch.nn.Module): + return next(obj.parameters()).device.type + elif isinstance(obj, torch.Tensor): + return "cuda" if obj.is_cuda else "cpu" + else: + raise ValueError("") + + +def check_iterable(val): + """duck-type check if val is iterable""" + try: + _ = iter(val) + except: + return False + else: + return True + + +def remove_by_idx(input, idx, dim): + """ + remove part of input indexed by idx on dim + """ + idx_keep = [i for i in range(input.size(dim)) if i not in idx] + return torch.index_select(input, dim, torch.tensor(idx_keep).to(input.device)) + + +def get_eval_transform_by_cutoff(cutoff): + """ + get method for transforming eigenvalues into a binary keep fraction + + will scale each eigenvector by 1 or 0 depending on whether that eigenvalue + explains more than **cutoff** fraction of the variance + + returns a callable method + """ + + def eval_transform(evals): + assert torch.all(evals >= 0), "found negative eigenvalues, doesn't work for 'cutoff' eval_transform" + evals = evals / torch.sum(evals) + return 1.0 * (evals > cutoff) + + return eval_transform + + +def fractional_histogram(*args, **kwargs): + """wrapper of np.histogram() with relative counts instead of total or density""" + counts, bins = np.histogram(*args, **kwargs) + counts = counts / np.sum(counts) + return counts, bins + + +def edge2center(edges): + """from a list of edges of bins (e.g. for torch.histogram()), return the centers between the edges""" + assert edges.ndim == 1, "edges must be a 1-d array" + return edges[:-1] + np.diff(edges) / 2 + + +def smartcorr(input): + """ + Performs torch corrcoef on the input data but sets each pair-wise correlation coefficent + to 0 where the activity has no variance (var=0) for a particular dimension (replaces nans with zeros)sss + """ + idx_zeros = torch.var(input, dim=1) == 0 + cc = torch.corrcoef(input) + cc[idx_zeros, :] = 0 + cc[:, idx_zeros] = 0 + return cc + + +def batch_cov(input, centered=True, correction=True): + """ + Performs batched covariance on input data of shape (batch, dim, samples) or (dim, samples) + + Where the resulting batch covariance matrix has shape (batch, dim, dim) or (dim, dim) + and bcov[i] = torch.cov(input[i]) if input.ndim==3 + + if centered=True (default) will subtract the means first + + if correction=True, will use */(N-1) otherwise will use */N + """ + assert (input.ndim == 2) or (input.ndim == 3), "input must be a 2D or 3D tensor" + assert isinstance(correction, bool), "correction must be a boolean variable" + + # check if batch dimension was provided + no_batch = input.ndim == 2 + + # add an empty batch dimension if not provided + if no_batch: + input = input.unsqueeze(0) + + # measure number of samples of each input matrix + S = input.size(2) + + # subtract mean if doing centered covariance + if centered: + input = input - input.mean(dim=2, keepdim=True) + + # measure covariance of each input matrix + bcov = torch.bmm(input, input.transpose(1, 2)) + + # correct for number of samples + bcov /= S - 1.0 * correction + + # remove empty batch dimension if not provided + if no_batch: + bcov = bcov.squeeze(0) + + return bcov + + +def smart_pca(input, centered=True, use_rank=True, correction=True): + """ + smart algorithm for pca optimized for speed + + input should either have shape (batch, dim, samples) or (dim, samples) + if dim > samples, will use svd and if samples < dim will use covariance/eigh method + + will center data when centered=True + + if it fails, will fall back on performing sklearns IncrementalPCA whenever forcetry=True + """ + assert (input.ndim == 2) or (input.ndim == 3), "input should be a matrix or batched matrices" + assert isinstance(correction, bool), "correction should be a boolean" + + if input.ndim == 2: + no_batch = True + input = input.unsqueeze(0) # create batch dimension for uniform code + else: + no_batch = False + + _, D, S = input.size() + if D > S: + # if more dimensions than samples, it's more efficient to run svd + v, w, _ = named_transpose([torch.linalg.svd(inp) for inp in input]) + # convert singular values to eigenvalues + w = [ww**2 / (S - 1.0 * correction) for ww in w] + # append zeros because svd returns w in R**k where k = min(D, S) + w = [torch.concatenate((ww, torch.zeros(D - S))) for ww in w] + + else: + # if more samples than dimensions, it's more efficient to run eigh + bcov = batch_cov(input, centered=centered, correction=correction) + w, v = named_transpose([eigendecomposition(C, use_rank=use_rank) for C in bcov]) + + # return to stacked tensor across batch dimension + w = torch.stack(w) + v = torch.stack(v) + + # if no batch originally provided, squeeze out batch dimension + if no_batch: + w = w.squeeze(0) + v = v.squeeze(0) + + # return eigenvalues and eigenvectors + return w, v + + +def eigendecomposition(C, use_rank=True): + """ + helper for getting eigenvalues and eigenvectors of covariance matrix + + will measure eigenvalues and eigenvectors with torch.linalg.eigh() + the output will be sorted from highest to lowest eigenvalue (& eigenvector) + + if use_rank=True, will measure the rank of the covariance matrix and zero + out any eigenvalues beyond the rank (that are usually nonzero numerical errors) + """ + try: + # measure eigenvalues and eigenvectors + w, v = torch.linalg.eigh(C) + + except torch._C._LinAlgError as error: + # this happens if the algorithm failed to converge + # try with sklearn's incrementalPCA algorithm + return sklearn_pca(C, use_rank=use_rank) + + except Exception as error: + # if any other exception, raise it + raise error + + # sort by eigenvalue from highest to lowest + w_idx = torch.argsort(-w) + w = w[w_idx] + v = v[:, w_idx] + + # iff use_rank=True, will set eigenvalues to 0 for probable numerical errors + if use_rank: + crank = torch.linalg.matrix_rank(C) # measure rank of covariance + w[crank:] = 0 # set eigenvalues beyond rank to 0 + + # return eigenvalues and eigenvectors + return w, v + + +def sklearn_pca(input, use_rank=True, rank=None): + """ + sklearn incrementalPCA algorithm serving as a replacement for eigh when it fails + + input should be a tensor with shape (num_samples, num_features) or it can be a + covariance matrix with (num_features, num_features) + + if use_rank=True, will set num_components to the rank of input and then fill out the + rest of the components with random orthogonal components in the null space of the true + components and set the eigenvalues to 0 + + if use_rank=False, will attempt to fit all the components + if rank is not None, will attempt to fit #=rank components without measuring the rank directly + (will ignore "rank" if use_rank=False) + + returns w, v where w is eigenvalues and v is eigenvectors sorted from highest to lowest + """ + # dimension + num_samples, num_features = input.shape + + # measure rank (or set to None) + rank = None if not use_rank else (rank if rank is not None else fast_rank(input)) + + # create and fit IncrementalPCA object on input data + ipca = IncrementalPCA(n_components=rank).fit(input) + + # eigenvectors are the components + v = ipca.components_ + + # eigenvalues are the scaled singular values + w = ipca.singular_values_**2 / num_samples + + # if v is a subspace of input (e.g. not a full basis, fill it out) + if v.shape[0] < num_features: + msg = "adding this because I think it should always be true, and if not I want to find out" + assert w.shape[0] == v.shape[0], msg + v_kernel = null_space(v).T + v = np.vstack((v, v_kernel)) + w = np.concatenate((w, np.zeros(v_kernel.shape[0]))) + + return torch.tensor(w, dtype=torch.float), torch.tensor(v, dtype=torch.float).T + + +def fast_rank(input): + """uses transpose to speed up rank computation, otherwise normal""" + if input.size(-2) < input.size(-1): + input = torch.transpose(input, -2, -1) + return int(torch.linalg.matrix_rank(input)) + + +# ------------------ alignment functions ---------------------- +def alignment(input, weight, method="alignment", relative=True): + """ + measure alignment (proportion variance explained) between **input** and **weight** + + computes the rayleigh quotient between each weight vector in **weight** and the **input** fed + into **weight**. Typically, **input** is the output in Layer L-1 and **weight** is from Layer L + + the output is normalized by the total variance in output of layer L-1 to measure the proportion + of variance of in **input** is explained by a projection onto node's weights in **weight** + + args + ---- + input: (batch, neurons) torch tensor + - represents input activity being fed in to network weight layer + weight: (num_out, num_in) torch tensor + - represents weights multiplied by input layer + method: string, default='alignment' + - which method to use to measure structure in **input** + - if 'alignment', uses covariance matrix of **input** + - if 'similarity', uses correlation matrix of **input** + relative: bool, default=True, + - if True, will measure relative RQ (divide by sum of eigenvalues) + + returns + ------- + alignment: (num_out, ) torch tensor + - proportion of variance explained by projection of **input** onto each **weight** vector + """ + assert method == "alignment" or method == "similarity", "method must be set to either 'alignment' or 'similarity' (or None, default is alignment)" + if method == "alignment": + cc = torch.cov(input.T) + #cc = torch.corrcoef(input.T) + elif method == "similarity": + cc = smartcorr(input.T) + else: + raise ValueError(f"did not recognize method ({method}), must be 'alignment' or 'similarity'") + + rq = torch.sum(torch.matmul(weight, cc) * weight, axis=1) / torch.sum(weight * weight, axis=1) + + if relative: + return rq / torch.trace(cc) + + # shuffled_rq = rq[torch.randperm(rq.size(0))] # Shuffle RQ values + # rq = shuffled_rq / torch.trace(cc) + + return rq + + +@torch.no_grad() +def expected_alignment_distribution(eigenvalues, relative=True, valid_rotation=True, with_rotation=True, bins=11, num_tests=100): + """ + for a set of eigenvalues, measure the expected distribution given aligned weights + + relative determines whether to normalize by sum of eigenvalues + valid_rotation determines whether we create orthonormal rotation matrices (for True) + or just normally distributed weights with the expected variance (for False) + bins works like histogram bins + num_tests determines how many tests to do (it's actually num_tests*len(eigenvalues)) + """ + # otherwise, randomly sample using eigenvalue as weighted average + N = len(eigenvalues) + if relative: + eigenvalues /= eigenvalues.sum() + eigenvalues = eigenvalues.view(-1, 1).expand(-1, N * num_tests) + if with_rotation: + if valid_rotation: + mixing = [torch.linalg.qr(torch.normal(0, 1 / math.sqrt(N), (N, N)))[0].T for _ in range(num_tests)] + coefficients = torch.concatenate(mixing, axis=1) ** 2 + else: + coefficients = torch.normal(0, 1 / math.sqrt(N), (N, N * num_tests)) ** 2 + else: + coefficients = torch.ones((N, N * num_tests)) + coefficients = coefficients.to(get_device(eigenvalues)) + weights = eigenvalues * coefficients + alignment = torch.sum(eigenvalues * weights, dim=0) / weights.sum(dim=0) + counts, bins = torch.histogram(alignment.cpu(), bins=bins, density=True) + centers = edge2center(bins) + return counts, bins, centers + +def compute_skewness_inplace(input): + """ + Compute skewness for each input feature (dimension) using in-place operations. + """ + mean = input.mean(dim=0) + variance = input.var(dim=0) + + # Add epsilon to variance to avoid division by zero + epsilon = 1e-6 + variance = variance + epsilon + + third_moment = (input - mean).pow(3).mean(dim=0) + + # Skewness formula: E[(X - μ)^3] / (E[(X - μ)^2])^(3/2) + skewness = third_moment / variance.pow(1.5) + + return skewness + +def compute_kurtosis_inplace(input): + """ + Compute kurtosis for each input feature (dimension) using in-place operations. + """ + mean = input.mean(dim=0) + variance = input.var(dim=0) + fourth_moment = (input - mean).pow_(4).mean(dim=0) + + # Kurtosis formula: E[(X - μ)^4] / (E[(X - μ)^2])^2 - 3 + kurtosis = fourth_moment / variance.pow(2) - 3 + return kurtosis + +def compute_skewness_low_rank(input, rank_approx=50): + """ + Compute skewness after reducing the dimensionality using PCA to the specified rank approximation. + """ + + pca = PCA(n_components=rank_approx) + input_reduced = torch.Tensor(pca.fit_transform(input.cpu().numpy())).to(input.device) + + return compute_skewness_inplace(input_reduced) + +def compute_kurtosis_low_rank(input, rank_approx=50): + """ + Compute kurtosis after reducing the dimensionality using PCA to the specified rank approximation. + """ + + pca = PCA(n_components=rank_approx) + input_reduced = torch.Tensor(pca.fit_transform(input.cpu().numpy())).to(input.device) + + return compute_kurtosis_inplace(input_reduced) + +def compute_redundancy(weights, input_covariance): + """ + Compute the redundancy matrix for all pairs of nodes using matrix multiplication. + + Args: + weights: (n, d) torch tensor, where n is the number of nodes and d is the input dimension. + input_covariance: (d, d) torch tensor, the covariance matrix of the input data. + + Returns: + redundancy_matrix: (n, n) torch tensor, redundancy between each pair of nodes (diagonal excluded). + """ + # Compute the redundancy matrix: R = W Σ_X W^T + redundancy_matrix = torch.matmul(weights, torch.matmul(input_covariance, weights.T)) + + # Zero out the diagonal elements (self-information) + redundancy_matrix.fill_diagonal_(0) + + return redundancy_matrix + + + +def alignment_expansion(input, weight, method="alignment_expansion", relative=True): + """ + Measure alignment (proportion variance explained) between **input** and **weight** + and compute single-node information, redundancy, and total information for the layer. + + args: + ---- + input: (batch, neurons) torch tensor + - represents input activity being fed into network weight layer + weight: (num_out, num_in) torch tensor + - represents weights multiplied by input layer + method: string, default='alignment' + - which method to use to measure structure in **input** + - if 'alignment', uses covariance matrix of **input** + - if 'similarity', uses correlation matrix of **input** + relative: bool, default=True, + - if True, will measure relative RQ (divide by sum of eigenvalues) + + returns: + ------- + alignment: (num_out, ) torch tensor + - proportion of variance explained by projection of **input** onto each **weight** vector + single_node_info: torch tensor + - Information carried by each node (mutual information for single nodes) + redundancy_matrix: torch tensor + - Pairwise redundancy (mutual information overlap between pairs of nodes) + total_info: float + - Total information of the layer, accounting for redundancy + """ + + # Step 1: Compute covariance of input + # if method == "alignment_expansion": + cc = torch.cov(input.T) # Covariance matrix of input + # elif method == "alignment_expansion": + # cc = torch.corrcoef(input.T) # Correlation matrix of input + # else: + # raise ValueError(f"Method {method} not recognized. Use 'alignment' or 'similarity'.") + + # Step 2: Compute Rayleigh Quotient (RQ) for each node + rq = torch.sum(torch.matmul(weight, cc) * weight, axis=1) / torch.sum(weight * weight, axis=1) + + # Step 3: Compute Single-Node Information for each node (with kurtosis correction) + single_node_info = torch.log(rq)# / torch.trace(cc) # First term (proportional to RQ) + + # Compute third-order term: Skewness correction + skewness = compute_skewness_inplace(input) + + # Ensure the skewness shape matches the input dimensionality + if skewness.shape[0] == input.shape[1]: # If skewness is (d,), reshape it + skewness = skewness.view(1, -1) # Reshape skewness to (1, d) for broadcasting + + # Apply the skewness correction term + normalized_weight = weight #/ weight.norm(p=2, dim=1, keepdim=True) + skewness_correction = torch.sum(normalized_weight * skewness, dim=1) + + # Compute kurtosis-based correction term + kurtosis = compute_kurtosis_low_rank(input) # Kurtosis is computed for each input feature + + # Expand kurtosis tensor to match weight shape + # if kurtosis.dim() == 1: + # kurtosis = kurtosis.unsqueeze(0) # Add a dimension to match the weight tensor shape + + kurtosis_correction = 0.5 * torch.norm(normalized_weight, p=2, dim=1) * kurtosis[:weight.shape[1]].mean() + + #single_node_info = skewness_correction + #single_node_info = kurtosis_correction + + # Step 4: Continue with redundancy and total information as before + redundancy_matrix = torch.abs(compute_redundancy(normalized_weight, cc)) + adjusted_single_node_info = adjust_information_with_redundancy(single_node_info, redundancy_matrix)#single_node_info.clone() + + total_info = torch.sum(adjusted_single_node_info) + + if method == "alignment_0": + info = adjusted_single_node_info + elif method == "alignment_1": + info = skewness_correction + elif method == "alignment_2": + info = kurtosis_correction + elif method == "alignment_red": + info = redundancy_matrix + + + return info + +def adjust_information_with_redundancy(single_node_info, redundancy_matrix): + """ + Adjust single-node information by subtracting redundancy for each pair of nodes. + + Args: + single_node_info: (n, ) torch tensor, the single-node information for each node. + redundancy_matrix: (n, n) torch tensor, the redundancy between each pair of nodes. + + Returns: + adjusted_info: (n, ) torch tensor, adjusted single-node information. + """ + n = single_node_info.size(0) + + # Compute pairwise differences between single-node information + #info_diffs = single_node_info.view(n, 1) - single_node_info.view(1, n) + + # Create a mask where the single-node information is smaller for each pair + #mask = (info_diffs < 0).float() + + # Subtract redundancy from the node with smaller information + #redundancy_adjustment = torch.sum(mask * redundancy_matrix, dim=1) + redundancy_matrix.fill_diagonal_(0) + + redundancy_adjustment = torch.sum(redundancy_matrix, dim=1) + + # Adjust the single-node information + adjusted_info =redundancy_adjustment + + return adjusted_info + +def plot_information_results(single_node_info, adjusted_single_node_info, redundancy_matrix, total_info): + """ + Plot single-node information, adjusted single-node information, redundancy matrix, and total information for the layer. + + args: + ---- + single_node_info: torch tensor + - Information carried by each node + adjusted_single_node_info: torch tensor + - Adjusted information for each node after subtracting redundancy + redundancy_matrix: torch tensor + - Pairwise redundancy between nodes + total_info: float + - Total information of the layer + """ + import matplotlib.pyplot as plt + import seaborn as sns + + # Plot single-node information + plt.figure(figsize=(8, 6)) + plt.bar(range(len(single_node_info)), single_node_info.cpu().detach().numpy(), alpha=0.6, label="Original Info") + plt.bar(range(len(adjusted_single_node_info)), adjusted_single_node_info.cpu().detach().numpy(), alpha=0.6, label="Adjusted Info") + plt.xlabel('Node Index') + plt.ylabel('Single Node Information (MI)') + plt.title('Single-Node Information for Each Node (Original vs Adjusted)') + plt.legend() + plt.show() + + # Plot redundancy matrix + plt.figure(figsize=(8, 6)) + sns.heatmap(redundancy_matrix.cpu().detach().numpy(), annot=True, cmap="YlGnBu") + plt.title('Redundancy Matrix (Pairwise Redundancy between Nodes)') + plt.show() + + # Display total information + print(f"Total Information for the Layer: {total_info.item():.4f}") + + + +def get_maximum_strides(h_input, w_input, layer): + h_max = int(np.floor((h_input + 2 * layer.padding[0] - layer.dilation[0] * (layer.kernel_size[0] - 1) - 1) / layer.stride[0] + 1)) + w_max = int(np.floor((w_input + 2 * layer.padding[1] - layer.dilation[1] * (layer.kernel_size[1] - 1) - 1) / layer.stride[1] + 1)) + return h_max, w_max + + +def get_unfold_params(layer): + return dict(stride=layer.stride, padding=layer.padding, dilation=layer.dilation) + + +# ----- cvPCA methods ----- +@torch.no_grad() +def cvPCA(X1, X2): + """X1, X2 are both (dimensions x samples)""" + D, B = X1.shape + assert X2.shape == (D, B), "shape of X1 and X2 is not the same" + _, u = smart_pca(X1) + + cproj0 = X1.T @ u + cproj1 = X2.T @ u + ss = (cproj0 * cproj1).mean(axis=0) + return ss + + +def get_num_components(nc, shape): + return nc if nc is not None else min(shape) + + +@torch.no_grad() +def shuff_cvPCA(X1, X2, nshuff=5, cvmethod=cvPCA): + """X1, X2 are both (dimensions x samples)""" + D, B = X1.shape + assert X2.shape == (D, B), "shape of X1 and X2 is not the same" + nc = get_num_components(None, (D, B)) + ss = torch.zeros((nshuff, nc)) + X = torch.stack((X1, X2)) + for k in range(nshuff): + iflip = 1 * (torch.rand(B) > 0.5) + X1c = torch.gather(X, 0, iflip.view(1, 1, -1).expand(1, D, -1)).squeeze(0) + X2c = torch.gather(X, 0, -(iflip - 1).view(1, 1, -1).expand(1, D, -1)).squeeze(0) + ss[k] = cvmethod(X1c, X2c) + return ss + + +def avg_value_by_layer(full): + """ + return average value per layer across training + + **full** is a list of lists where the outer list is each snapshot through training or + minibatch etc and each inner list is the value for each node in the network across layers + of a particular measurement + + For example: + num_epochs = 1000 + nodes_per_layer = [50, 40, 30, 20] + len(full) == 1000 + len(full[i]) == 4 ... for all i + [f.shape for f in full[i]] = [50, 40, 30, 20] ... for all i + + this method will return a tensor of size (num_layers, num_epochs) of the average value (for + whatever value is in **full**) for each list/list + """ + num_epochs = len(full) + num_layers = len(full[0]) + avg_full = torch.zeros((num_layers, num_epochs)) + for layer in range(num_layers): + avg_full[layer, :] = torch.tensor([torch.mean(f[layer]) for f in full]) + return avg_full.cpu() + + +def value_by_layer(full: List[List[torch.Tensor]], layer: int) -> torch.Tensor: + """ + return all value measurements for a particular layer from **full** + + **full** is a list of lists where the outer list is each snapshot through training or + minibatch etc and each inner list is the value for each node in the network across layers + + this method will return just the part of **full** corresponding to the layer indexed + by **layer** as a tensor of shape (num_epochs, num_nodes) + + see ``avg_value_by_layer`` for a little more explanation + """ + return torch.cat([f[layer].view(1, -1) for f in full], dim=0).cpu() + + +def condense_values(full: List[List[List[torch.Tensor]]]) -> List[torch.Tensor]: + """ + condense List[List[List[Tensor]]] representing some value measured across networks, batches, and layers, for each node in the layer + + returns list of #=num_layers tensors, where each tensor has shape (num_networks, num_batches, num_nodes_per_layer) + + full should be a list of list of lists + the first list should have length = number of networks + the second list should have length = number of batches + the third list should have length = number of layers in the network (this has to be the same for each network!) + the tensor should have shape = number of nodes in this layer (also must be the same for each network) (or can be anything as long as consistent across layers) + """ + num_layers = len(full[0][0]) + return [torch.stack([value_by_layer(value, layer) for value in full]) for layer in range(num_layers)] + + +def transpose_list(list_of_lists): + """helper function for transposing the order of a list of lists""" + return list(map(list, zip(*list_of_lists))) + + +def named_transpose(list_of_lists, reduction=None): + """ + helper function for transposing lists without forcing the output to be a list like transpose_list + + for example, if list_of_lists contains 10 copies of lists that each have 3 iterable elements you + want to name "A", "B", and "C", then write: + A, B, C = named_transpose(list_of_lists) + + if reduction is used, it will be applied to each output, otherwise will make them lists + """ + if reduction is not None: + return map(reduction, zip(*list_of_lists)) + return map(list, zip(*list_of_lists)) + + +def ptp(tensor, dim=None, keepdim=False): + """ + simple method for measuring range of tensor on requested dimension or on all data + """ + if dim is None: + return tensor.max() - tensor.min() + return tensor.max(dim, keepdim).values - tensor.min(dim, keepdim).values + + +def rms(tensor, dim=None, keepdim=False): + """simple method for measuring root-mean-square on requested dimension or on all data in tensor""" + if dim is None: + return torch.sqrt(torch.mean(tensor**2)) + return torch.sqrt(torch.mean(tensor**2, dim=dim, keepdim=keepdim)) + + +def compute_stats_by_type(tensor, num_types, dim, method="var"): + """ + helper method for returning the mean and variance across a certain dimension + where multiple types are concatenated on that dimension + + for example, suppose we trained 2 networks each with 3 sets of parameters + and concatenated the loss in a tensor like [set1-loss-net1, set1-loss-net2, set2-loss-net1, ...] + then this would contract across the nets from each set and return the mean and variance + """ + num_on_dim = tensor.size(dim) + num_per_type = int(num_on_dim / num_types) + tensor_by_type = tensor.unsqueeze(dim) + expand_shape = list(tensor_by_type.shape) + expand_shape[dim + 1] = num_per_type + expand_shape[dim] = num_types + tensor_by_type = tensor_by_type.view(expand_shape) + type_means = torch.mean(tensor_by_type, dim=dim + 1) + if method == "var": + type_dev = torch.var(tensor_by_type, dim=dim + 1) + elif method == "std": + type_dev = torch.std(tensor_by_type, dim=dim + 1) + elif method == "se": + type_dev = torch.std(tensor_by_type, dim=dim + 1) / np.sqrt(num_per_type) + elif method == "range": + type_dev = ptp(tensor_by_type, dim=dim + 1) + else: + raise ValueError(f"Method ({method}) not recognized.") + + return type_means, type_dev + + +def weighted_average(data, weights, dim, keepdim=False, ignore_nan=False): + """ + take the weighted average of **data** on a certain dimension with **weights** + + weights should be a nonnegative vector that broadcasts into data + avg = sum_i(data_i * weight_i, dim) / sum_i(weight_i, dim) + + if ignore_nan=True, (default=False), will ignore nans in weighted average + """ + assert data.ndim == weights.ndim, "data and weights must have same number of dimensions" + assert torch.all(weights[~torch.isnan(weights)] >= 0), "weights must be nonnegative" + + for d in dim if check_iterable(dim) else [dim]: + assert data.size(d) == weights.size( + d + ), f"data and weights must have same size in averaging dimensions (data.size({d})={data.size(d)}, (weight.size({d})={weights.size(d)}))" + + # use normal sum if not ignore nan, otherwise use nansum + sum = torch.nansum if ignore_nan else torch.sum + + # make sure nans are in the same place in weights and data for accurate division by total weight + if ignore_nan: + weights = weights.expand(data.size()) + weights = torch.masked_fill(weights, torch.isnan(data), torch.nan) + + # numerator & denominator of weighted average + numerator = sum(data * weights, dim=dim, keepdim=keepdim) + denominator = sum(weights, dim=dim, keepdim=keepdim) + + # return weighted average + return numerator / denominator + + +def fgsm_attack(image, epsilon, data_grad, transform, sign): + """update an image with fast-gradient sign method""" + warn("fgsm_attack is only going to be in utils temporarily!", DeprecationWarning, stacklevel=2) + # Collect the element-wise sign of the data gradient + if sign: + data_grad = data_grad.sign() + else: + data_grad = data_grad.clone() + # Create the perturbed image by adjusting each pixel of the input image + perturbed_image = image + epsilon * data_grad + # Adding clipping to maintain [0,1] range + perturbed_image = transform(perturbed_image) + # Return the perturbed image + return perturbed_image + + +def str2bool(str): + if isinstance(str, bool): + return str + if str.lower() in ("true", "1"): + return True + elif str.lower() in ("false", "0"): + return False + else: + raise TypeError("Boolean type expected") + + +def save_checkpoint(nets, optimizers, results, path): + """ + Method for saving checkpoints for networks throughout training. + """ + multi_model_ckpt = {f"model_state_dict_{i}": net.state_dict() for i, net in enumerate(nets)} + multi_optimizer_ckpt = {f"optimizer_state_dict_{i}": opt.state_dict() for i, opt in enumerate(optimizers)} + checkpoint = results | multi_model_ckpt | multi_optimizer_ckpt + torch.save(checkpoint, path) + + +def load_checkpoints(nets, optimizers, device, path): + """ + Method for loading presaved checkpoint during training. + TODO: device handling for passing between gpu/cpu + """ + + if device == "cpu": + checkpoint = torch.load(path, map_location=device) + elif device == "cuda": + checkpoint = torch.load(path) + + net_ids = natsorted([key for key in checkpoint if key.startswith("model_state_dict")]) + opt_ids = natsorted([key for key in checkpoint if key.startswith("optimizer_state_dict")]) + assert all( + [oi.split("_")[-1] == ni.split("_")[-1] for oi, ni in zip(opt_ids, net_ids)] + ), "nets and optimizers cannot be matched up from checkpoint" + + [net.load_state_dict(checkpoint.pop(net_id)) for net, net_id in zip(nets, net_ids)] + [opt.load_state_dict(checkpoint.pop(opt_id)) for opt, opt_id in zip(optimizers, opt_ids)] + + if device == "cuda": + [net.to(device) for net in nets] + + return nets, optimizers, checkpoint + + +def match_git(path): + """simple method for determining if a path is a git-related file or directory""" + if ".git" in path: + return True + return False + + +def compress_directory(output_path, directory_path=None): + """send an entire directory to a zip file at output_path, using .gitignore and ignoring .git files""" + if directory_path is None: + # relative to utils -- this is the main repo path + directory_path = os.path.dirname(os.path.abspath(__file__)) + "/../.." + print(directory_path) + # Parse .gitignore file + gitignore_path = os.path.join(directory_path, ".gitignore") + matches = parse_gitignore(gitignore_path) + + # Prepare list for copying files + files_to_copy = [] + archive_names = [] + for dirpath, dirnames, files in os.walk(directory_path): + if matches(dirpath) or match_git(dirpath): + # clear any files from within this path + dirnames[:] = [] + else: + # Filter files based on .gitignore rules (and don't save any .git files) + keep_files = [f for f in files if not matches(f) and not match_git(f)] + # Make full path + full_files = [os.path.join(dirpath, f) for f in keep_files] + for file in full_files: + # Add file to the copy list + files_to_copy.append(file) + archive_names.append(os.path.relpath(file, directory_path)) + + # create zip file + with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf: + # go through directory + for file, name in zip(files_to_copy, archive_names): + zipf.write(file, arcname=name) diff --git a/_arxiv/_archive/benchmark_dropout.py b/_arxiv/_archive/benchmark_dropout.py new file mode 100644 index 00000000..1bd94dad --- /dev/null +++ b/_arxiv/_archive/benchmark_dropout.py @@ -0,0 +1,230 @@ +""" +Benchmark script for comparing different dropout implementation approaches. + +This script compares the performance of three approaches to progressive dropout: +1. Original (sequential): Processing networks one-by-one +2. Batched: Processing networks in small batches +3. Tensorized: Using tensor operations to process networks at once +""" + +import os +import sys +import time +import argparse +import logging +from typing import Dict, List, Tuple, Any +from collections import defaultdict + +import numpy as np +import torch +import torch.nn as nn +from tqdm import tqdm + +from alignment.config import ExperimentConfig +from alignment.models.registry import create_model +from alignment.metrics import get_metric +from alignment.datasets import load_dataset +from alignment.dropout import progressive_dropout +from alignment import utils + +# Setup logging +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") + +def create_networks(config, num_networks=20): + """Create multiple neural networks for benchmarking.""" + networks = [] + device = config.device + + for i in range(num_networks): + # Set different seed for each network + if hasattr(config, 'seed') and config.seed is not None: + torch.manual_seed(config.seed + i) + torch.cuda.manual_seed_all(config.seed + i) + np.random.seed(config.seed + i) + + # Create model + model = create_model(config.model) + model.to(device) + networks.append(model) + + logger.info(f"Created {len(networks)} models with independent initializations") + return networks + +def run_benchmark(config_path, approaches=["original", "batched", "tensorized"], + num_networks=20, num_runs=3, num_batches=5, detailed_timing=True): + """ + Run benchmarks for different dropout implementation approaches. + + Args: + config_path: Path to configuration file + approaches: List of approaches to benchmark + num_networks: Number of networks to use + num_runs: Number of times to run each benchmark for averaging + num_batches: Number of data batches to process + detailed_timing: Whether to time individual steps + """ + # Load configuration + config = ExperimentConfig.load(config_path) + + # Create networks + networks = create_networks(config, num_networks) + + # Prepare dataset + dataset = load_dataset(config.dataset) + + # Get metric + metric = get_metric(config.alignment.metric) + + # Benchmark results + results = {} + detailed_results = defaultdict(lambda: defaultdict(list)) + + # Run benchmarks + for approach in approaches: + logger.info(f"Benchmarking {approach} approach") + + # Set approach-specific parameters + if approach == "original": + # Original sequential approach + kwargs = {"network_batch_size": 1, "use_tensorized": False} + elif approach == "batched": + # Batched approach + kwargs = {"network_batch_size": 4, "use_tensorized": False} + elif approach == "tensorized": + # Tensorized approach + kwargs = {"use_tensorized": True} + else: + logger.error(f"Unknown approach: {approach}") + continue + + # Run multiple times and average + timings = [] + for run in range(num_runs): + logger.info(f"Run {run+1}/{num_runs}") + + # Time the execution of full dropout + start_time = time.time() + + # Run progressive dropout with the specified approach + dropout_results = progressive_dropout( + networks, + dataset, + dropout_fractions=np.linspace(0.1, 0.9, config.alignment.dropout_steps), + metric=metric, + device=config.device, + pruning_mode=config.extra.dropout_pruning_mode, + dropout_mode=config.extra.dropout_mode, + num_batches=num_batches, + detailed_timing=detailed_timing, + **kwargs + ) + + # Calculate elapsed time + elapsed_time = time.time() - start_time + timings.append(elapsed_time) + + # Collect detailed timing information if available + if detailed_timing and hasattr(dropout_results, "timing_info"): + for key, value in dropout_results.timing_info.items(): + detailed_results[approach][key].append(value) + + # Log detailed timing information + for key, values in dropout_results.timing_info.items(): + if isinstance(values, (int, float)): + logger.info(f" {key}: {values:.2f} seconds") + + logger.info(f"Completed run in {elapsed_time:.2f} seconds") + + # Calculate average and standard deviation + avg_time = np.mean(timings) + std_time = np.std(timings) + + # Store results + results[approach] = { + "average_time": avg_time, + "std_time": std_time, + "timings": timings + } + + logger.info(f"{approach}: {avg_time:.2f} ± {std_time:.2f} seconds") + + # Print comparison + logger.info("\nBenchmarking Results:") + logger.info("-" * 50) + logger.info(f"{'Approach':<15} {'Average Time (s)':<20} {'Speedup':<10}") + logger.info("-" * 50) + + # Calculate speedup relative to original approach + original_time = results["original"]["average_time"] if "original" in results else None + + for approach in approaches: + if approach not in results: + continue + + avg_time = results[approach]["average_time"] + speedup = original_time / avg_time if original_time else 1.0 + + logger.info(f"{approach:<15} {avg_time:.2f} ± {results[approach]['std_time']:.2f} {'':>5} {speedup:.2f}x") + + logger.info("-" * 50) + + # Print detailed timing if available + if detailed_timing: + logger.info("\nDetailed Timing Breakdown:") + logger.info("-" * 80) + + # Find all unique timing keys + all_keys = set() + for approach in approaches: + if approach in detailed_results: + all_keys.update(detailed_results[approach].keys()) + + # Print header + logger.info(f"{'Operation':<25} " + " ".join([f"{approach:<15}" for approach in approaches])) + logger.info("-" * 80) + + # Print timings for each key + for key in sorted(all_keys): + line = f"{key:<25} " + for approach in approaches: + if approach in detailed_results and key in detailed_results[approach]: + values = detailed_results[approach][key] + if len(values) > 0: + avg = np.mean(values) + line += f"{avg:.2f}s{' ' * 10}" + else: + line += f"{'N/A':<15}" + else: + line += f"{'N/A':<15}" + logger.info(line) + + logger.info("-" * 80) + + return results, detailed_results if detailed_timing else None + +def main(): + """Main function for benchmarking.""" + parser = argparse.ArgumentParser(description="Benchmark dropout implementation approaches") + parser.add_argument("--config", type=str, required=True, help="Path to configuration file") + parser.add_argument("--approaches", type=str, nargs="+", default=["original", "batched", "tensorized"], + help="Approaches to benchmark") + parser.add_argument("--num-networks", type=int, default=20, help="Number of networks to use") + parser.add_argument("--num-runs", type=int, default=3, help="Number of runs for each benchmark") + parser.add_argument("--num-batches", type=int, default=5, help="Number of data batches to process") + parser.add_argument("--no-detailed-timing", action="store_true", help="Disable detailed timing") + + args = parser.parse_args() + + # Run benchmarks + run_benchmark( + args.config, + args.approaches, + args.num_networks, + args.num_runs, + args.num_batches, + not args.no_detailed_timing + ) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/_arxiv/_archive/config_alignment_stats.yaml b/_arxiv/_archive/config_alignment_stats.yaml new file mode 100644 index 00000000..abd73f95 --- /dev/null +++ b/_arxiv/_archive/config_alignment_stats.yaml @@ -0,0 +1,56 @@ +experiment: alignment_stats +results_path: # <-- Add path to where results to be saved/loaded +no_save: False +just_plot: False +save_networks: False +show_params: False +show_all: False +#device: cpu +use_timestamp: False +#timestamp: # The timestamp of previous experiment to plot. + +dataset: + name: MNIST + path: # <-- Add path to your dataset + download: True # Downloads it if it does not exist. + +model: + name: MLP + dropout: 0 + alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"} + +optimizer: + name: Adam + lr: 1e-3 + weight_decay: 0 + +training: + batch_size: 1024 + epochs: 1 + replicates: 1 + +alignment: + no_alignment: False + delta_weights: False + frequency: 1 + +checkpointing: + use_prev: False + save_checkpoints: False + frequency: 1 + use_wandb: False + +extra: + num_drops: 9 + dropout_by_layer: False + + ### For alignment_comparison and loading_prediction experiments + comparison: "lr" + regularizers: ["none", "dropout", "weight_decay"] + lrs: [1e-2, 1e-3, 1e-4] + compare_dropout: 0.5 + compare_wd: 1e-5 + + ### For adversarial_shaping experiment + cutoffs: [1e-2, 1e-3, 1e-4, 0.0] + manual_frequency: 5 diff --git a/_arxiv/_archive/config_alignment_stats_old.yaml b/_arxiv/_archive/config_alignment_stats_old.yaml new file mode 100644 index 00000000..a25906ba --- /dev/null +++ b/_arxiv/_archive/config_alignment_stats_old.yaml @@ -0,0 +1,66 @@ +experiment: alignment_stats +results_path: results # <-- Add path to where results to be saved/loaded +no_save: False +just_plot: False +save_networks: False +show_params: False +show_all: False +device: cuda +use_timestamp: False +#timestamp: # The timestamp of previous experiment to plot. + +dataset: + name: MNIST + path: /n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data + +model: + name: MLP + dropout: 0 + alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"} + +training: + batch_size: 1024 + epochs: 3 + replicates: 3 + name: Adam + lr: 1e-3 + weight_decay: 0 + +alignment: + do_alignment: True + frequency: 1 + methods: ["RQ"] + measure_expected: False + bins: 50 + cnn_mode: "unfold" # Options for CNN alignment processing: + # "unfold" - Reshapes conv layer features into (-1, features) with batch and patches together (default) + # "patchwise" - Computes alignment per-patch, with special processing in patchwise_alignment + # "batch_patch_combined" - Standardized approach that combines batch and patch dimensions + # Differences between modes: + # - "unfold" is the classic approach with flexible reshaping options + # - "patchwise" preserves patch structure for specialized calculations + # - "batch_patch_combined" ensures consistent dimensionality across frameworks --> original method + scale_by_norm: False # If True, scales covariance matrices by their norm before computing RQ + # Similar to the "similarity" option in original codebase +checkpointing: + use_prev: False + save_checkpoints: False + frequency: 1 + use_wandb: True + +extra: + num_drops: 40 + dropout_pruning_mode: "global_joint" # Options: + # "global_joint" - Prune X% of nodes across all layers based on their global alignment ranking + # This distributes pruning across layers proportional to where the + # highest/lowest alignment nodes are located. Some layers may have more + # neurons pruned than others, preserving the global ranking of importance. + # "layer_wise" - Prune exactly X% from each layer and apply pruning to all layers at once + # This shows the effect of pruning all layers simultaneously with fixed proportions + # "layer_isolated" - Prune X% from only one layer at a time + # This creates multiple figures (one per layer) showing isolated effects + dropout_mode: "scaled" # Options for how targeted dropout is applied to neurons: + # "scaled" - Zeroes neurons and applies scaling to maintain signal magnitude (default) + # "unscaled" - Zeroes neurons but doesn't apply compensation scaling + exclude_classification_layer: True # If True, excludes the classification layer from pruning (always true in original codebase) + \ No newline at end of file diff --git a/_arxiv/_archive/config_alignment_stats_v2.yaml b/_arxiv/_archive/config_alignment_stats_v2.yaml new file mode 100644 index 00000000..f1e50278 --- /dev/null +++ b/_arxiv/_archive/config_alignment_stats_v2.yaml @@ -0,0 +1,74 @@ +experiment: "mnist_alignment_example" +results_path: # <-- Add path to where results to be saved/loaded +no_save: false ## if true, skip saving results and plots +just_plot: false ## If true, skip the experiment’s main() logic, only load previously saved results and call plot_from_existing() +save_networks: false ## If true, also save trained networks (self.save_networks(...)) after training. +show_params: false ## keep the plots open WE CAN REMOVE THIS +show_all: false ## ## keep the plots open WE CAN REMOVE THIS +use_timestamp: false ## If use_timestamp is true, the code stores results in subfolders named by timestamp (an ISO/time string). +timestamp: null + +ddp_world_size: 2 + +dataset: + name: "ImageNet" # "MNIST", "CIFAR10", "CIFAR100", "ImageNet", etc. + path: # <-- Add path to your dataset ## path to dataset + +model: + name: "AlexNet" ## base artitecture, e.g., "MLP", "CNN2P2", "AlexNet". + dropout: 0 ## Dropout fraction 0 to 1 + alignment_layers: null + # "layer2.0": "layer1.0" + # a dict of layer names: source_layer. if null it will be set to all layers. + #"features.8": null # the null after "features.8" refers to the fact that we don't mention a particular input for this layer and use what it gets + +training: + do_train: true + epochs: 10 + batch_size: 100 + lr: 1e-3 + replicates: 2 + name: "Adam" + weight_decay: 0.0 + +alignment: +# do_train_alignment: false +# do_test_alignment: true + do_alignment: false + methods: ["RQ"] + measure_expected: False + frequency: 1 + bins: 50 + cnn_mode: "unfold" #"unfold: Single covariance across all patches", "patch: Patchwise alignment, weighting by variance", "old"" + +checkpointing: + use_wandb: true + use_prev: false + save_checkpoints: false + frequency: 1 + +plots: + show_loss: true + show_dropout: true + show_eigenfeatures: true + show_eig_dropout: true + +extra: + num_drops: 50 + dropout_by_layer: false + single_layer_mode: false # dropout_by_layer must be true to work + progressive_dropout_on_train: false + aggregate_alignment: true + +# • If aggregate_alignment=True, we directly append the alignment each batch to results["alignment"]. +# • If aggregate_alignment=False, we temporarily store in epoch_align_data and only append once at the end of the epoch. + + #comparison: "lr" ## for “alignment_comparison” style experiments + #regularizers: ["none", "dropout", "weight_decay"] + #lrs: [1e-3, 1e-4] + #compare_dropout: 0.5 + #compare_wd: 1e-5 + #cutoffs: [1e-2, 1e-3, 1e-4, 0.0] + #manual_frequency: 5 + + # python src/alignment_v2/experiments/alignment_experiments.py --config .vscode/config_alignment_stats_v2.yaml \ No newline at end of file diff --git a/_arxiv/_archive/config_pruning_modes.yaml b/_arxiv/_archive/config_pruning_modes.yaml new file mode 100644 index 00000000..9e5fb726 --- /dev/null +++ b/_arxiv/_archive/config_pruning_modes.yaml @@ -0,0 +1,80 @@ +# Pruning Modes Configuration Sample + +# This config demonstrates the different pruning modes available: +# 1. "global_joint": Concatenate all nodes and prune X% across all layers (original v1) +# 2. "layer_wise": Prune X% from each layer separately, but all at once (v2-like) +# 3. "layer_isolated": Prune X% from only one layer at a time (new option) +# 4. "cascading_layer": Prune layers progressively - prune layer 1, use pruned network for layer 2, etc. + +experiment: + name: pruning_modes_comparison + type: alignment_stats + +dataset: + name: MNIST + batch_size: 200 + test_batch_size: 200 + +model: + network: MLP + input_dim: 784 + hidden_dim: [512, 256, 128, 64] + output_dim: 10 + activation: ReLU + dropout: 0.2 + +alignment: + methods: + - "RQ" + frequency: 10 + measure_expected: true + do_alignment: true + bins: 21 + +training: + epochs: 50 + lr: 0.001 + optimizer: Adam + criterion: CrossEntropyLoss + device: cuda:0 + +checkpointing: + use_prev: false + save_checkpoints: false + frequency: 10 + +extra: + # Choose one of the pruning modes: + + # Option 1: Global Joint Pruning (Original v1) + # Concatenates all nodes from all layers and prunes the lowest X% across all nodes + # dropout_pruning_mode: "global_joint" + + # Option 2: Layer-wise Pruning (v2-like) + # Prunes X% from each layer separately, but all at once + # dropout_pruning_mode: "layer_wise" + + # Option 3: Layer Isolation Pruning (new) + # Prunes X% from only one layer at a time, generating separate results per layer + # dropout_pruning_mode: "layer_isolated" + + # Option 4: Cascading Layer Pruning (new) + # Progressively prunes layers - prune layer 1, use pruned network to compute alignments for layer 2, etc. + dropout_pruning_mode: "cascading_layer" + + aggregate_alignment: false + num_drops: 9 + +# Alternative configs you can uncomment to use: + +# Config for "layer_wise" mode (v2 behavior) +# extra: +# dropout_pruning_mode: "layer_wise" +# aggregate_alignment: false +# num_drops: 9 + +# Config for "layer_isolated" mode +# extra: +# dropout_pruning_mode: "layer_isolated" +# aggregate_alignment: false +# num_drops: 9 \ No newline at end of file diff --git a/_arxiv/_archive/config_refactored_test.yaml b/_arxiv/_archive/config_refactored_test.yaml new file mode 100644 index 00000000..5c409d7e --- /dev/null +++ b/_arxiv/_archive/config_refactored_test.yaml @@ -0,0 +1,65 @@ +experiment: "alignment_experiment" +results_path: "results/pruning_test" +no_save: false +just_plot: false +save_networks: true +show_params: false +show_all: false +device: "cuda" +use_timestamp: true +timestamp: null +ddp_world_size: 1 + +dataset: + dataset_name: "MNIST" + dataset_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data" + +model: + model_name: "mlp" + dropout: 0.0 + alignment_layers: null + +training: + do_train: true + epochs: 5 + batch_size: 1024 + replicates: 1 + name: "Adam" + lr: 1e-3 + weight_decay: 0.0 + +alignment: + do_alignment: true + metric: "RQ" + methods: ["RQ"] + measure_expected: false + frequency: 1 + bins: 50 + cnn_mode: "unfold" + scale_by_norm: false + run_progressive: true + run_eigenvector: true + dropout_min: 0.0 + dropout_max: 0.9 + dropout_steps: 10 + +checkpointing: + use_prev: false + save_checkpoints: true + frequency: 1 + use_wandb: false + +plots: + show_loss: true + show_dropout: true + show_eigenfeatures: true + show_eig_dropout: true + +extra: + num_drops: 10 + progressive_dropout_on_train: false + single_layer_mode: false + aggregate_alignment: false + dropout_pruning_mode: "per_layer_combined" + dropout_mode: "scaled" + exclude_classification_layer: true \ No newline at end of file diff --git a/_arxiv/_archive/config_with_eigenvector.yaml b/_arxiv/_archive/config_with_eigenvector.yaml new file mode 100644 index 00000000..5da61835 --- /dev/null +++ b/_arxiv/_archive/config_with_eigenvector.yaml @@ -0,0 +1,59 @@ +# Alignment Experiment Configuration +# This configuration file defines parameters for neural network alignment experiments +# that analyze the relationship between alignment, dropout, and network performance. + +# Experiment metadata +experiment_type: "alignment" # Type of experiment to run (options: "alignment", "training", "analysis") +results_path: "results" # Base path for saving results (can be relative or absolute path) +use_timestamp: true # Whether to include timestamp in result paths for unique experiment tracking + +# Core execution options +device: "cuda" # Device to run on (options: "cuda", "cpu", or specific GPU like "cuda:0") +no_save: false # If true, don't save any results to disk (useful for debug runs) +just_plot: false # If true, only load and plot existing results without running experiments +save_networks: true # Whether to save trained network checkpoints for later use +show_all: false # Whether to display all plots during execution (vs. just saving them) + +# Dataset configuration +dataset: + dataset_name: "MNIST" # Name of the dataset to use (options: "MNIST", "CIFAR10", "CIFAR100") + data_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data" # Path to dataset files + batch_size: 1024 # Batch size for training and evaluation (higher values use more memory but train faster) + +# Model configuration +model: + model_name: "MLP" # Model architecture to use (options: "MLP", "CNN2P2", "alexnet") + dropout_rate: 0.0 # Base dropout rate during training (range: 0.0-1.0, 0.0 means no dropout) + alignment_layers: null # Layers to measure alignment on (null for all layers, or specify as [0,1,2]) + +# Training configuration +training: + epochs: 5 # Number of training epochs (more epochs can improve performance but take longer) + replicates: 10 # Number of network replicates to train (for statistical analysis) + optimizer: "Adam" # Optimizer to use (options: "Adam", "SGD", "RMSprop") + learning_rate: 1.0e-3 # Learning rate for optimizer (typical range: 1e-4 to 1e-2) + weight_decay: 0.0 # L2 regularization strength (higher values prevent overfitting) + +# Checkpointing configuration +checkpointing: + save_checkpoints: true # Whether to save training checkpoints to disk + checkpoint_frequency: 1 # How often to save checkpoints (in epochs, 1 means every epoch) + use_wandb: true # Whether to use Weights & Biases for experiment tracking and visualization + load_checkpoint: false # Whether to load from previous checkpoint (for resuming training) + +# Alignment analysis configuration +alignment: + metric: "RQ" # Alignment metric to use (options: "RQ", "NullSpace", "Delta", "R2") + run_progressive: true # Whether to run progressive dropout experiment (node pruning) + run_eigenvector: true # Whether to run eigenvector dropout experiment (feature pruning based on PCA components) + dropout_min: 0.0 # Minimum dropout fraction to test (range: 0.0-1.0) + dropout_max: 0.95 # Maximum dropout fraction to test (range: 0.0-1.0, <1.0 to avoid complete dropout) + dropout_steps: 20 # Number of dropout fractions to test (more steps = higher resolution but slower) + scale_by_norm: false # Whether to scale covariance matrices by norm (can improve numerics) + +# Extra configuration parameters +extra: + dropout_mode: "scaled" # How dropout is applied (options: "scaled" = multiply by 1/(1-p), "unscaled" = no scaling) + dropout_pruning_mode: "global_joint" # How pruning is distributed (options: "global_joint" = across all neurons, "layer_wise" = layer-by-layer, "layer_isolated" = one layer at a time) + exclude_classification_layer: true # Whether to exclude the output layer from pruning (usually true for classification tasks) + cnn_mode: "unfold" # How CNN features are processed for alignment (options: "unfold", "patchwise", "batch_patch_combined") \ No newline at end of file diff --git a/_arxiv/_archive/configs/config_alignment_experiment.yaml b/_arxiv/_archive/configs/config_alignment_experiment.yaml new file mode 100644 index 00000000..1850e894 --- /dev/null +++ b/_arxiv/_archive/configs/config_alignment_experiment.yaml @@ -0,0 +1,203 @@ +# Alignment Experiment Configuration +# This configuration file defines parameters for neural network alignment experiments +# that analyze the relationship between alignment, dropout, and network performance. + +# ----------------------------------------------------------------------------- +# I. Experiment Metadata & Core Execution +# ----------------------------------------------------------------------------- +experiment_name: "alexnet_imagenet_alignment_test" +# experiment_name: "cnn_mnist_nested_config_test" +experiment_type: "AUTO" # Options: + # "AUTO" - Automatically selects the best experiment type based on the dataset. + # "progressive_dropout" - Prunes nodes iteratively and measures performance. + # "eigenvector_dropout" - Prunes based on eigenvectors of alignment metric. + # "alignment_analysis" - Runs both progressive and (if configured) eigenvector dropout. + # "layer_isolated_pruning" - Prunes/evaluates one layer at a time (generates N plots). + # "cascading_layer_pruning" - Prunes layers sequentially, L1 then L2 etc (simple version). +results_path: "results" # Base path for saving all outputs. +use_timestamp: true # If true, appends a YYYYMMDD_HHMMSS timestamp to `results_path`. +debug_mode: false # MODIFIED: Set to true for testing callbacks + +device: "cuda" # "cuda", "cuda:0", "cpu". Ignored if use_ddp is true (device set by local_rank). +save_networks: true # If true, final trained/pruned network models are saved. +no_save: false # Master switch: if true, nothing is saved to disk. +just_plot: false # If true, loads existing results and only generates plots. +show_all: false # If true, matplotlib plots will be displayed interactively. +seed: 422 # MODIFIED: Set a seed for reproducibility during testing + +# --- NEW: DDP (Distributed Data Parallel) Configuration --- +use_ddp: false # Set to true to enable DDP for multi-GPU training. + # Requires launching with torchrun or similar. + # e.g., torchrun --nproc_per_node=NUM_GPUS your_script.py --config this_file.yaml +ddp_backend: "nccl" # DDP backend. "nccl" is recommended for NVIDIA GPUs. +# ddp_rank, ddp_world_size, ddp_local_rank are set automatically by the launch script if use_ddp is true. +# --- End DDP Configuration --- + +# ----------------------------------------------------------------------------- +# II. Dataset Configuration +# ----------------------------------------------------------------------------- +# Maps to DatasetConfig in config.py +dataset: + dataset_name: "ImageNet" # Options: "MNIST", "CIFAR10", "CIFAR100", "ImageNet". + data_path: "/n/holylfs06/LABS/kempner_shared/Lab/data/imagenet_1k/" # Root directory for dataset files. + batch_size: 256 # Batch size for DataLoaders (reduced for ImageNet due to larger images). + num_workers: 8 # Number of worker processes for DataLoader (increased for ImageNet). + +# ----------------------------------------------------------------------------- +# III. Model Configuration (NEW STRUCTURE) +# ----------------------------------------------------------------------------- +# Maps to ModelConfig in config.py +model: + # --- Identifier and Common Parameters --- + model_name: "torchvision_alexnet" # Options: "MLP", "CNN2P2", "torchvision_alexnet", "torchvision_resnet18", "hf_bert-base-uncased", etc. + output_dim: 1000 # Number of output classes. Used by the final layer of the model. + dropout_rate: 0.0 # General dropout rate that internal model constructors might use. + + # --- CNN Mode for AlignmentNetwork Wrapper --- + # Defines how the AlignmentNetwork wrapper (if the model is a CNN) preprocesses its inputs for metric calculation. + # Options: "unfold", "patchwise", "batch_patch_combined". Default in code is "unfold". + cnn_mode: "unfold" + + # --- Alignment Layers Specification (Common) --- + # Optional: Specify layers for AlignmentNetwork to analyze. Names must match model.named_modules(). + # If null, AlignmentNetwork attempts to find all layers with weights. + alignment_layers: null + # Example for AlexNet: + # alignment_layers: + # "features.0": 0 + # "features.3": 1 + # "classifier.6": 5 + + # --- Model-Specific Parameter Blocks --- + # Only include and populate ONE of the following blocks based on model_name. + + # (Example 1) MLP Parameters (if model_name: "MLP") + mlp_params: + input_dim: 784 + hidden_dims: [100, 100, 50] # Example: Two hidden layers + activation: "relu" + + # CNN2P2 Parameters (if model_name: "CNN2P2") - Comment out if using MLP or external + cnn2p2_params: + in_channels: 1 + conv_channels: [32, 64] + kernel_sizes: [5, 5] + strides: [1, 1] + paddings: [2, 2] + pool_kernel_size: 2 + pool_stride: 2 + hidden_fc_dim: 128 + example_input_hw: [28, 28] + + # (Example 3) External Model Parameters (if model_name indicates an external model, e.g., "torchvision_alexnet") + external_params: # Used if model_name is "torchvision_alexnet", "hf_bert-base-uncased", or just "external" + source: "torchvision" # Options: "torchvision", "huggingface_transformers" + name_or_path: "alexnet" # e.g., "alexnet", "resnet18" for torchvision; "bert-base-uncased" for hf + pretrained: true # Load pretrained weights (True/False) + freeze_feature_extractor: false # Freeze all layers except the classifier (True/False) + + # --- Ad-hoc Parameters for Internal Constructors (Use Sparingly) --- + extra_model_params: {} # For any other specific kwargs to pass to internal model constructors (create_mlp, create_cnn2p2) + +# ----------------------------------------------------------------------------- +# IV. Training Configuration +# ----------------------------------------------------------------------------- +# Maps to TrainingConfig in config.py +training: + epochs: 5 + replicates: 1 # For DDP, typically train 1 replicate, DDP handles data parallelism. + # If not DDP, can set >1 for multiple independent training runs. + optimizer: "Adam" + learning_rate: 0.001 + momentum: 0.9 + loss: "cross_entropy" + training_method: "auto" # "auto", "sequential", "fully_tensorized". + # "fully_tensorized" is for non-DDP multi-replicate speedup on one GPU. + # With DDP, "sequential" is typically used for the single replicate. + train_before_dropout: false # Whether to run initial training phase before dropout experiments. + +# ----------------------------------------------------------------------------- +# V. Alignment Metric Calculation Settings +# ----------------------------------------------------------------------------- +# Maps to alignment_settings: AlignmentConfig in config.py +alignment_settings: + metric: + - "RQ" + #- "MI_G" + #- "Node_Redundancy" + scale_by_norm: false # Global default for scale_by_norm + + # cnn_mode below is for how metric calculation functions (e.g. compute_all_node_scores) + # process CNN features IF they receive them in a specific format (e.g. from AlignmentNetwork's precomputation). + # The ModelConfig.cnn_mode determines how AlignmentNetwork itself preprocesses and provides these features. + # These two should ideally be consistent or ModelConfig.cnn_mode dictates the format received by metric functions. + cnn_mode: "unfold" # Options: "unfold", "patchwise", "batch_patch_combined", "filter_patch_summary", "filter_specific_covariance_rq" + # Detailed options for alignment_settings.cnn_mode (how metrics process CNN features): + # "unfold": Assumes features are [batch_size * num_patches, feature_dim] (e.g., from AlignmentNetwork cnn_mode="unfold"). + # For RQ, computes ONE covariance matrix from ALL these patch features (from ONE original data batch). + # Each filter's RQ is then computed against this single, global covariance matrix. + # Matches typical RQ definition but can be memory-intensive. + # "filter_patch_summary": (Specific to RQ metric) For each filter, computes a score based on its interaction + # with every input patch it sees (e.g., mean/max of SQUARED COSINE SIMILARITIES + # between filter and patches). Memory-efficient. Not mathematically identical to covariance-based RQ. + # "filter_specific_covariance_rq": (Specific to RQ metric) For each filter, computes its full covariance-based RQ. + # The covariance is derived from all unfolded input patches (from ONE original batch) + # that the filter processes. Mathematically standard RQ per filter. Computationally intensive. + # "patchwise": Expects features like [batch_size, num_patches, feature_dim]. Intended for metrics that process per-patch. + # (Currently requires metric-specific support not fully integrated into general path for all metrics). + # "batch_patch_combined": Similar effect to "unfold" for how data is shaped for global covariance calculations. + + cnn_rq_aggregation_op: "mean" # If cnn_mode for metric is "filter_patch_summary" for RQ: "mean", "max", "var", "sum". + force_cpu_for_large_metric_ops: false # If true, attempts to offload large tensor ops in metrics (like cov mat) to CPU. + + run_progressive: true # For experiment_type="alignment_analysis", whether to run progressive dropout part. + run_eigenvector: false # For experiment_type="alignment_analysis", whether to run eigenvector dropout part. + + callbacks: # Configure metrics to track during training + alignment_metrics: + - name: "RQ" + num_batches: 5 # Batches for callback metrics. null for all test batches. + - name: "MI_G" # Gaussian MI + num_batches: 5 + # - name: "PID_SI" # PID Shared Info (Requires BROJA_2PID to be correctly placed and importable) + # num_batches: 1 + +# ----------------------------------------------------------------------------- +# VI. Pruning Experiment Settings +# ----------------------------------------------------------------------------- +# Maps to pruning_settings: PruningConfig in config.py +pruning_settings: + dropout_min: 0.0 + dropout_max: 1.0 + dropout_steps: 30 # Reduced for faster testing + + dropout_mode: "scaled" # "scaled" or "unscaled" + # "scaled": Remaining weights are scaled by 1/(1-p) (Inverted Dropout principle). + # "unscaled" (or "zero"): Weights are simply zeroed out. + dropout_pruning_mode: "layer_isolated" # "global_joint", "layer_wise", "layer_isolated", "cascading_layer" + # "global_joint": Considers all nodes from all prunable layers together. Ranks globally by score and prunes top X% overall. + # "layer_wise": Prunes X% from *each* prunable layer independently, based on scores within that layer. + # "layer_isolated": (For experiment_type="layer_isolated_pruning") Prunes X% in one layer, evaluates, restores, then moves to next layer. + # "cascading_layer": Prunes L1, its output influences pruning of L2, etc. (Note: current batched impl. has limitations). + exclude_classification_layer: true + use_multi_strategy_dropout: true + num_batches_for_scores: null # Number of batches to compute pruning scores on. + +# ----------------------------------------------------------------------------- +# VII. Checkpointing & Logging Configuration +# ----------------------------------------------------------------------------- +checkpointing: + save_checkpoints: false + checkpoint_frequency: 1 + load_checkpoint: false + +wandb: + use_wandb: true + wandb_project: "neural_alignment" + wandb_entity: "your_wandb_entity" # <<< IMPORTANT: Update with your W&B entity or remove line for default + +extra: + log_frequency: 1 + log_images: true + detailed_logging: true + # dummy_extra_param: null \ No newline at end of file diff --git a/_arxiv/_archive/configs/config_alignment_stats.yaml b/_arxiv/_archive/configs/config_alignment_stats.yaml new file mode 100644 index 00000000..abd73f95 --- /dev/null +++ b/_arxiv/_archive/configs/config_alignment_stats.yaml @@ -0,0 +1,56 @@ +experiment: alignment_stats +results_path: # <-- Add path to where results to be saved/loaded +no_save: False +just_plot: False +save_networks: False +show_params: False +show_all: False +#device: cpu +use_timestamp: False +#timestamp: # The timestamp of previous experiment to plot. + +dataset: + name: MNIST + path: # <-- Add path to your dataset + download: True # Downloads it if it does not exist. + +model: + name: MLP + dropout: 0 + alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"} + +optimizer: + name: Adam + lr: 1e-3 + weight_decay: 0 + +training: + batch_size: 1024 + epochs: 1 + replicates: 1 + +alignment: + no_alignment: False + delta_weights: False + frequency: 1 + +checkpointing: + use_prev: False + save_checkpoints: False + frequency: 1 + use_wandb: False + +extra: + num_drops: 9 + dropout_by_layer: False + + ### For alignment_comparison and loading_prediction experiments + comparison: "lr" + regularizers: ["none", "dropout", "weight_decay"] + lrs: [1e-2, 1e-3, 1e-4] + compare_dropout: 0.5 + compare_wd: 1e-5 + + ### For adversarial_shaping experiment + cutoffs: [1e-2, 1e-3, 1e-4, 0.0] + manual_frequency: 5 diff --git a/_arxiv/_archive/configs/config_alignment_stats_old.yaml b/_arxiv/_archive/configs/config_alignment_stats_old.yaml new file mode 100644 index 00000000..a25906ba --- /dev/null +++ b/_arxiv/_archive/configs/config_alignment_stats_old.yaml @@ -0,0 +1,66 @@ +experiment: alignment_stats +results_path: results # <-- Add path to where results to be saved/loaded +no_save: False +just_plot: False +save_networks: False +show_params: False +show_all: False +device: cuda +use_timestamp: False +#timestamp: # The timestamp of previous experiment to plot. + +dataset: + name: MNIST + path: /n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data + +model: + name: MLP + dropout: 0 + alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"} + +training: + batch_size: 1024 + epochs: 3 + replicates: 3 + name: Adam + lr: 1e-3 + weight_decay: 0 + +alignment: + do_alignment: True + frequency: 1 + methods: ["RQ"] + measure_expected: False + bins: 50 + cnn_mode: "unfold" # Options for CNN alignment processing: + # "unfold" - Reshapes conv layer features into (-1, features) with batch and patches together (default) + # "patchwise" - Computes alignment per-patch, with special processing in patchwise_alignment + # "batch_patch_combined" - Standardized approach that combines batch and patch dimensions + # Differences between modes: + # - "unfold" is the classic approach with flexible reshaping options + # - "patchwise" preserves patch structure for specialized calculations + # - "batch_patch_combined" ensures consistent dimensionality across frameworks --> original method + scale_by_norm: False # If True, scales covariance matrices by their norm before computing RQ + # Similar to the "similarity" option in original codebase +checkpointing: + use_prev: False + save_checkpoints: False + frequency: 1 + use_wandb: True + +extra: + num_drops: 40 + dropout_pruning_mode: "global_joint" # Options: + # "global_joint" - Prune X% of nodes across all layers based on their global alignment ranking + # This distributes pruning across layers proportional to where the + # highest/lowest alignment nodes are located. Some layers may have more + # neurons pruned than others, preserving the global ranking of importance. + # "layer_wise" - Prune exactly X% from each layer and apply pruning to all layers at once + # This shows the effect of pruning all layers simultaneously with fixed proportions + # "layer_isolated" - Prune X% from only one layer at a time + # This creates multiple figures (one per layer) showing isolated effects + dropout_mode: "scaled" # Options for how targeted dropout is applied to neurons: + # "scaled" - Zeroes neurons and applies scaling to maintain signal magnitude (default) + # "unscaled" - Zeroes neurons but doesn't apply compensation scaling + exclude_classification_layer: True # If True, excludes the classification layer from pruning (always true in original codebase) + \ No newline at end of file diff --git a/_arxiv/_archive/configs/config_alignment_stats_v2.yaml b/_arxiv/_archive/configs/config_alignment_stats_v2.yaml new file mode 100644 index 00000000..f1e50278 --- /dev/null +++ b/_arxiv/_archive/configs/config_alignment_stats_v2.yaml @@ -0,0 +1,74 @@ +experiment: "mnist_alignment_example" +results_path: # <-- Add path to where results to be saved/loaded +no_save: false ## if true, skip saving results and plots +just_plot: false ## If true, skip the experiment’s main() logic, only load previously saved results and call plot_from_existing() +save_networks: false ## If true, also save trained networks (self.save_networks(...)) after training. +show_params: false ## keep the plots open WE CAN REMOVE THIS +show_all: false ## ## keep the plots open WE CAN REMOVE THIS +use_timestamp: false ## If use_timestamp is true, the code stores results in subfolders named by timestamp (an ISO/time string). +timestamp: null + +ddp_world_size: 2 + +dataset: + name: "ImageNet" # "MNIST", "CIFAR10", "CIFAR100", "ImageNet", etc. + path: # <-- Add path to your dataset ## path to dataset + +model: + name: "AlexNet" ## base artitecture, e.g., "MLP", "CNN2P2", "AlexNet". + dropout: 0 ## Dropout fraction 0 to 1 + alignment_layers: null + # "layer2.0": "layer1.0" + # a dict of layer names: source_layer. if null it will be set to all layers. + #"features.8": null # the null after "features.8" refers to the fact that we don't mention a particular input for this layer and use what it gets + +training: + do_train: true + epochs: 10 + batch_size: 100 + lr: 1e-3 + replicates: 2 + name: "Adam" + weight_decay: 0.0 + +alignment: +# do_train_alignment: false +# do_test_alignment: true + do_alignment: false + methods: ["RQ"] + measure_expected: False + frequency: 1 + bins: 50 + cnn_mode: "unfold" #"unfold: Single covariance across all patches", "patch: Patchwise alignment, weighting by variance", "old"" + +checkpointing: + use_wandb: true + use_prev: false + save_checkpoints: false + frequency: 1 + +plots: + show_loss: true + show_dropout: true + show_eigenfeatures: true + show_eig_dropout: true + +extra: + num_drops: 50 + dropout_by_layer: false + single_layer_mode: false # dropout_by_layer must be true to work + progressive_dropout_on_train: false + aggregate_alignment: true + +# • If aggregate_alignment=True, we directly append the alignment each batch to results["alignment"]. +# • If aggregate_alignment=False, we temporarily store in epoch_align_data and only append once at the end of the epoch. + + #comparison: "lr" ## for “alignment_comparison” style experiments + #regularizers: ["none", "dropout", "weight_decay"] + #lrs: [1e-3, 1e-4] + #compare_dropout: 0.5 + #compare_wd: 1e-5 + #cutoffs: [1e-2, 1e-3, 1e-4, 0.0] + #manual_frequency: 5 + + # python src/alignment_v2/experiments/alignment_experiments.py --config .vscode/config_alignment_stats_v2.yaml \ No newline at end of file diff --git a/_arxiv/_archive/configs/config_pruning_modes.yaml b/_arxiv/_archive/configs/config_pruning_modes.yaml new file mode 100644 index 00000000..9e5fb726 --- /dev/null +++ b/_arxiv/_archive/configs/config_pruning_modes.yaml @@ -0,0 +1,80 @@ +# Pruning Modes Configuration Sample + +# This config demonstrates the different pruning modes available: +# 1. "global_joint": Concatenate all nodes and prune X% across all layers (original v1) +# 2. "layer_wise": Prune X% from each layer separately, but all at once (v2-like) +# 3. "layer_isolated": Prune X% from only one layer at a time (new option) +# 4. "cascading_layer": Prune layers progressively - prune layer 1, use pruned network for layer 2, etc. + +experiment: + name: pruning_modes_comparison + type: alignment_stats + +dataset: + name: MNIST + batch_size: 200 + test_batch_size: 200 + +model: + network: MLP + input_dim: 784 + hidden_dim: [512, 256, 128, 64] + output_dim: 10 + activation: ReLU + dropout: 0.2 + +alignment: + methods: + - "RQ" + frequency: 10 + measure_expected: true + do_alignment: true + bins: 21 + +training: + epochs: 50 + lr: 0.001 + optimizer: Adam + criterion: CrossEntropyLoss + device: cuda:0 + +checkpointing: + use_prev: false + save_checkpoints: false + frequency: 10 + +extra: + # Choose one of the pruning modes: + + # Option 1: Global Joint Pruning (Original v1) + # Concatenates all nodes from all layers and prunes the lowest X% across all nodes + # dropout_pruning_mode: "global_joint" + + # Option 2: Layer-wise Pruning (v2-like) + # Prunes X% from each layer separately, but all at once + # dropout_pruning_mode: "layer_wise" + + # Option 3: Layer Isolation Pruning (new) + # Prunes X% from only one layer at a time, generating separate results per layer + # dropout_pruning_mode: "layer_isolated" + + # Option 4: Cascading Layer Pruning (new) + # Progressively prunes layers - prune layer 1, use pruned network to compute alignments for layer 2, etc. + dropout_pruning_mode: "cascading_layer" + + aggregate_alignment: false + num_drops: 9 + +# Alternative configs you can uncomment to use: + +# Config for "layer_wise" mode (v2 behavior) +# extra: +# dropout_pruning_mode: "layer_wise" +# aggregate_alignment: false +# num_drops: 9 + +# Config for "layer_isolated" mode +# extra: +# dropout_pruning_mode: "layer_isolated" +# aggregate_alignment: false +# num_drops: 9 \ No newline at end of file diff --git a/_arxiv/_archive/configs/config_refactored_test.yaml b/_arxiv/_archive/configs/config_refactored_test.yaml new file mode 100644 index 00000000..5c409d7e --- /dev/null +++ b/_arxiv/_archive/configs/config_refactored_test.yaml @@ -0,0 +1,65 @@ +experiment: "alignment_experiment" +results_path: "results/pruning_test" +no_save: false +just_plot: false +save_networks: true +show_params: false +show_all: false +device: "cuda" +use_timestamp: true +timestamp: null +ddp_world_size: 1 + +dataset: + dataset_name: "MNIST" + dataset_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data" + +model: + model_name: "mlp" + dropout: 0.0 + alignment_layers: null + +training: + do_train: true + epochs: 5 + batch_size: 1024 + replicates: 1 + name: "Adam" + lr: 1e-3 + weight_decay: 0.0 + +alignment: + do_alignment: true + metric: "RQ" + methods: ["RQ"] + measure_expected: false + frequency: 1 + bins: 50 + cnn_mode: "unfold" + scale_by_norm: false + run_progressive: true + run_eigenvector: true + dropout_min: 0.0 + dropout_max: 0.9 + dropout_steps: 10 + +checkpointing: + use_prev: false + save_checkpoints: true + frequency: 1 + use_wandb: false + +plots: + show_loss: true + show_dropout: true + show_eigenfeatures: true + show_eig_dropout: true + +extra: + num_drops: 10 + progressive_dropout_on_train: false + single_layer_mode: false + aggregate_alignment: false + dropout_pruning_mode: "per_layer_combined" + dropout_mode: "scaled" + exclude_classification_layer: true \ No newline at end of file diff --git a/_arxiv/_archive/configs/config_with_eigenvector.yaml b/_arxiv/_archive/configs/config_with_eigenvector.yaml new file mode 100644 index 00000000..5da61835 --- /dev/null +++ b/_arxiv/_archive/configs/config_with_eigenvector.yaml @@ -0,0 +1,59 @@ +# Alignment Experiment Configuration +# This configuration file defines parameters for neural network alignment experiments +# that analyze the relationship between alignment, dropout, and network performance. + +# Experiment metadata +experiment_type: "alignment" # Type of experiment to run (options: "alignment", "training", "analysis") +results_path: "results" # Base path for saving results (can be relative or absolute path) +use_timestamp: true # Whether to include timestamp in result paths for unique experiment tracking + +# Core execution options +device: "cuda" # Device to run on (options: "cuda", "cpu", or specific GPU like "cuda:0") +no_save: false # If true, don't save any results to disk (useful for debug runs) +just_plot: false # If true, only load and plot existing results without running experiments +save_networks: true # Whether to save trained network checkpoints for later use +show_all: false # Whether to display all plots during execution (vs. just saving them) + +# Dataset configuration +dataset: + dataset_name: "MNIST" # Name of the dataset to use (options: "MNIST", "CIFAR10", "CIFAR100") + data_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data" # Path to dataset files + batch_size: 1024 # Batch size for training and evaluation (higher values use more memory but train faster) + +# Model configuration +model: + model_name: "MLP" # Model architecture to use (options: "MLP", "CNN2P2", "alexnet") + dropout_rate: 0.0 # Base dropout rate during training (range: 0.0-1.0, 0.0 means no dropout) + alignment_layers: null # Layers to measure alignment on (null for all layers, or specify as [0,1,2]) + +# Training configuration +training: + epochs: 5 # Number of training epochs (more epochs can improve performance but take longer) + replicates: 10 # Number of network replicates to train (for statistical analysis) + optimizer: "Adam" # Optimizer to use (options: "Adam", "SGD", "RMSprop") + learning_rate: 1.0e-3 # Learning rate for optimizer (typical range: 1e-4 to 1e-2) + weight_decay: 0.0 # L2 regularization strength (higher values prevent overfitting) + +# Checkpointing configuration +checkpointing: + save_checkpoints: true # Whether to save training checkpoints to disk + checkpoint_frequency: 1 # How often to save checkpoints (in epochs, 1 means every epoch) + use_wandb: true # Whether to use Weights & Biases for experiment tracking and visualization + load_checkpoint: false # Whether to load from previous checkpoint (for resuming training) + +# Alignment analysis configuration +alignment: + metric: "RQ" # Alignment metric to use (options: "RQ", "NullSpace", "Delta", "R2") + run_progressive: true # Whether to run progressive dropout experiment (node pruning) + run_eigenvector: true # Whether to run eigenvector dropout experiment (feature pruning based on PCA components) + dropout_min: 0.0 # Minimum dropout fraction to test (range: 0.0-1.0) + dropout_max: 0.95 # Maximum dropout fraction to test (range: 0.0-1.0, <1.0 to avoid complete dropout) + dropout_steps: 20 # Number of dropout fractions to test (more steps = higher resolution but slower) + scale_by_norm: false # Whether to scale covariance matrices by norm (can improve numerics) + +# Extra configuration parameters +extra: + dropout_mode: "scaled" # How dropout is applied (options: "scaled" = multiply by 1/(1-p), "unscaled" = no scaling) + dropout_pruning_mode: "global_joint" # How pruning is distributed (options: "global_joint" = across all neurons, "layer_wise" = layer-by-layer, "layer_isolated" = one layer at a time) + exclude_classification_layer: true # Whether to exclude the output layer from pruning (usually true for classification tasks) + cnn_mode: "unfold" # How CNN features are processed for alignment (options: "unfold", "patchwise", "batch_patch_combined") \ No newline at end of file diff --git a/_arxiv/_archive/configs/training_example.yaml b/_arxiv/_archive/configs/training_example.yaml new file mode 100644 index 00000000..886566d1 --- /dev/null +++ b/_arxiv/_archive/configs/training_example.yaml @@ -0,0 +1,47 @@ +# Example configuration with tensorized training options +experiment_type: "progressive_dropout" +experiment_name: "Tensorized Training Example" +use_timestamp: true +device: "cuda" + +# Dataset configuration +dataset: + dataset_name: "MNIST" + batch_size: 128 + data_path: "./data" + +# Model configuration +model: + model_name: "MLP" + dropout_rate: 0.0 + +# Training configuration +training: + epochs: 5 + replicates: 10 # Number of network copies to train + learning_rate: 0.001 + weight_decay: 0.0 + optimizer: "Adam" + train_before_dropout: true + +# Alignment configuration +alignment: + metric: "RQ" + run_progressive: true + run_eigenvector: false + dropout_min: 0.0 + dropout_max: 0.95 + dropout_steps: 20 + +# Extra configuration options +extra: + dropout_mode: "scaled" + dropout_pruning_mode: "global_joint" + exclude_classification_layer: true + # Available training methods: "auto", "sequential", "tensorized", "fully_tensorized" + training_method: "fully_tensorized" # Using the fastest method + +# Checkpoint configuration +checkpointing: + save_checkpoints: false + use_wandb: false \ No newline at end of file diff --git a/_arxiv/_archive/debug_pruning_strategies.py b/_arxiv/_archive/debug_pruning_strategies.py new file mode 100644 index 00000000..3eb729d7 --- /dev/null +++ b/_arxiv/_archive/debug_pruning_strategies.py @@ -0,0 +1,235 @@ +import torch +import torch.nn as nn +import numpy as np +import matplotlib.pyplot as plt +import logging +import os +import sys +import copy + +# Add parent directory to path to access src +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '.'))) + +from src.alignment.models.models import MLP +from src.alignment.datasets import load_dataset +from src.alignment.dropout import test_pruning_strategies + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s [%(levelname)s] %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', +) +logger = logging.getLogger(__name__) + +def main(): + # Set random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Determine device + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + logger.info(f"Using device: {device}") + + # Load MNIST dataset for testing + dataset_config = { + "dataset_name": "MNIST", + "data_path": "data", + "batch_size": 1024, + "model_name": "mlp", + } + dataset = load_dataset(dataset_config, device=device) + logger.info(f"Loaded {dataset_config['dataset_name']} dataset") + + # Create a simple MLP model for MNIST + network = MLP( + input_dim=784, # 28x28 for MNIST + num_hidden=[100, 100], + output_dim=10, + dropout_rate=0.0, + ) + + # Get linear layers from the MLP model + # First find all linear layers in the network + linear_layers = [] + for name, module in network.named_modules(): + if isinstance(module, nn.Linear): + linear_layers.append(module) + + # Set alignment layers for pruning (all linear layers) + network.alignment_layers = linear_layers + network.alignment_layer_names = [f"Layer {i}" for i in range(len(network.alignment_layers))] + + network.to(device) + logger.info(f"Created MLP model with {len(network.alignment_layers)} alignment layers") + + # Train the network for a single epoch + logger.info("Training network...") + train_network(network, dataset, device) + + # Print statistics about the original model + logger.info("Evaluating original model...") + print_layer_stats(network) + orig_accuracy, orig_loss = dataset.evaluate(network, device) + logger.info(f"Original model - Accuracy: {orig_accuracy:.2f}%, Loss: {orig_loss:.4f}") + + # Save original weights for verification + original_weights = {} + for i, layer in enumerate(network.alignment_layers): + if hasattr(layer, "weight") and layer.weight is not None: + original_weights[i] = layer.weight.data.clone() + + # Test different pruning strategies + logger.info("Testing pruning strategies...") + pruning_percents = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9] + + # Test each strategy individually to confirm they differ + strategies = ["high_rq", "low_rq", "random"] + all_results = {} + + for strategy in strategies: + logger.info(f"Testing {strategy} pruning strategy...") + + # Make a fresh copy of the network for each strategy + network_copy = copy.deepcopy(network).to(device) + + # Verify the copy has the same weights + for i, layer in enumerate(network_copy.alignment_layers): + if i in original_weights: + weight_diff = torch.sum(torch.abs(layer.weight.data - original_weights[i])).item() + if weight_diff > 1e-5: + logger.error(f"Weights for layer {i} changed before pruning! Diff: {weight_diff}") + + # Test this specific strategy + strategy_results = test_pruning_strategies( + network_copy, + dataset, + pruning_percents=pruning_percents, + device=device, + strategy=strategy + ) + + # Store results for this strategy + all_results[strategy] = strategy_results + + # Print results + logger.info(f"{strategy} pruning results:") + for i, percent in enumerate(pruning_percents): + acc = strategy_results[f"{strategy}_acc"][i] + loss = strategy_results[f"{strategy}_loss"][i] + logger.info(f" {percent*100:.0f}% pruned: Accuracy = {acc:.2f}%, Loss = {loss:.4f}") + + # Plot results + plt.figure(figsize=(12, 6)) + for strategy in strategies: + results = all_results[strategy] + plt.plot( + [p*100 for p in pruning_percents], + results[f"{strategy}_acc"], + marker="o" if strategy == "high_rq" else ("s" if strategy == "low_rq" else "^"), + label=f"{strategy}", + linewidth=2 + ) + + plt.xlabel("Pruning Percentage", fontsize=12) + plt.ylabel("Accuracy (%)", fontsize=12) + plt.title("Pruning Strategies Comparison", fontsize=14) + plt.grid(True, linestyle="--", alpha=0.7) + plt.legend(fontsize=11) + plt.tight_layout() + + # Save plot + plt.savefig("pruning_strategies_debug.png", dpi=300, bbox_inches="tight") + logger.info("Saved plot to pruning_strategies_debug.png") + + # Check if strategies are different + if len(strategies) > 1: + are_different = False + for i in range(len(strategies)-1): + for j in range(i+1, len(strategies)): + s1, s2 = strategies[i], strategies[j] + accs1 = all_results[s1][f"{s1}_acc"] + accs2 = all_results[s2][f"{s2}_acc"] + + # Calculate difference + diff = np.mean(np.abs(np.array(accs1) - np.array(accs2))) + logger.info(f"Mean accuracy difference between {s1} and {s2}: {diff:.2f}%") + + if diff > 1.0: # Consider different if average difference > 1% + are_different = True + + if are_different: + logger.info("SUCCESS: Pruning strategies show different results!") + else: + logger.error("FAILURE: Pruning strategies show very similar results.") + +def train_network(network, dataset, device, num_epochs=3): + """Train a network for a specified number of epochs""" + logger.info(f"Training a single network") + optimizer = torch.optim.Adam(network.parameters(), lr=0.001) + + # Training loop + for epoch in range(num_epochs): + running_loss = 0.0 + correct = 0 + total = 0 + + # Create DataLoader for training + train_loader = torch.utils.data.DataLoader( + dataset.train_dataset, batch_size=1024, shuffle=True + ) + + # Train for one epoch + network.train() + for i, (inputs, targets) in enumerate(train_loader): + # Move data to device + inputs, targets = inputs.to(device), targets.to(device) + + # Zero the parameter gradients + optimizer.zero_grad() + + # Forward pass + outputs = network(inputs) + loss = torch.nn.functional.cross_entropy(outputs, targets) + + # Backward pass and optimize + loss.backward() + optimizer.step() + + # Calculate accuracy + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + # Update running loss + running_loss += loss.item() + + # Print epoch statistics + epoch_loss = running_loss / len(train_loader) + epoch_acc = 100. * correct / total + logger.info(f"Epoch {epoch+1}/{num_epochs}: Loss={epoch_loss:.4f}, Acc={epoch_acc:.2f}%") + + # Evaluate on test set + test_acc, test_loss = dataset.evaluate(network, device) + logger.info(f"Epoch {epoch+1}/{num_epochs}: Train Loss={epoch_loss:.4f}, Train Acc={epoch_acc:.2f}%, Test Loss={test_loss:.4f}, Test Acc={test_acc:.2f}%") + +def print_layer_stats(network): + """Print statistics about each alignment layer in the network""" + for i, layer in enumerate(network.alignment_layers): + if hasattr(layer, "weight") and layer.weight is not None: + weights = layer.weight.data + total_weights = weights.numel() + zero_weights = (weights == 0).sum().item() + zero_percent = 100.0 * zero_weights / total_weights + + # Count pruned neurons (rows with all zeros) + pruned_neurons = 0 + total_neurons = weights.size(0) + for j in range(total_neurons): + if torch.all(weights[j] == 0): + pruned_neurons += 1 + + logger.info(f"Layer {i}: Shape {weights.shape}, zeros: {zero_weights}/{total_weights} ({zero_percent:.2f}%), pruned neurons: {pruned_neurons}/{total_neurons} ({100.0 * pruned_neurons / total_neurons:.2f}%)") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/_arxiv/_archive/md_files/DOCUMENTATION.md b/_arxiv/_archive/md_files/DOCUMENTATION.md new file mode 100644 index 00000000..01a39254 --- /dev/null +++ b/_arxiv/_archive/md_files/DOCUMENTATION.md @@ -0,0 +1,52 @@ +# Network Alignment Analysis Documentation + +This document provides a comprehensive index to all documentation for the Network Alignment Analysis codebase. + +## Getting Started + +- [README](README.md) - Main repository README with setup instructions +- [Usage Guide](doc/usage.md) - How to use the codebase +- [Installation Guide](README.md#installation) - How to install the package + +## Core Documentation + +- [Main Documentation](doc/DOCUMENTATION.md) - Overview of codebase architecture and capabilities +- [API Reference](doc/api/README.md) - Comprehensive API reference +- [Metrics System](doc/metrics/README.md) - Documentation for the metrics system +- [Experiment Framework](doc/experiment/README.md) - Documentation for the experiment framework +- [Performance Optimizations](doc/performance/README.md) - Documentation for performance optimizations +- [Tensorized Implementations](doc/tensorized/README.md) - Documentation for tensorized implementations + +## Guides and References + +- [Pruning Modes](doc/pruning_modes.md) - Documentation for different pruning strategies +- [Background](doc/background.md) - Background information on alignment analysis +- [Development Roadmap](doc/ROADMAP.md) - Future development plans + +## Refactoring and Cleanup + +- [Metrics Refactoring Summary](METRICS_REFACTORING_SUMMARY.md) - Summary of metrics system refactoring +- [Codebase Cleanup Summary](CODEBASE_CLEANUP_SUMMARY.md) - Summary of codebase cleanup + +## Directory Structure + +The codebase is organized into the following directories: + +- `src/alignment/`: Core source code implementing alignment metrics and algorithms + - `metrics.py`: Implementation of all alignment metrics + - `utils/`: Utility functions for data handling, visualization, etc. + - `experiment/`: Experiment framework for running alignment experiments + - `networks/`: Network architectures and training utilities + - `pruning/`: Implementation of various pruning strategies + +- `tests/`: Unit and integration tests +- `scripts/`: Utility scripts for running experiments and analysis +- `benchmarks/`: Performance evaluation scripts +- `configs/`: Configuration files for experiments +- `results/`: Output directory for experiment results +- `doc/`: Documentation + +## Additional Resources + +- [CHANGELOG.md](CHANGELOG.md) - History of changes to the codebase +- [LICENSE](LICENSE) - MIT License for the project \ No newline at end of file diff --git a/_arxiv/_archive/md_files/README_tensorized_dropout.md b/_arxiv/_archive/md_files/README_tensorized_dropout.md new file mode 100644 index 00000000..4daf37ff --- /dev/null +++ b/_arxiv/_archive/md_files/README_tensorized_dropout.md @@ -0,0 +1,93 @@ +# Tensorized Dropout and Training Documentation + +This document explains the tensorized implementations for training and dropout operations in the alignment project. + +## Tensorized Training Methods + +The project now includes optimized tensorized training methods for efficiently training multiple networks in parallel. This is particularly useful for experiments that require training multiple networks with the same architecture but different initializations. + +### Available Training Methods + +1. **Sequential Training** (`sequential`): Original method that trains each network one at a time. +2. **Tensorized Training** (`tensorized`): Trains networks in parallel by batching their training steps. +3. **Fully Tensorized Training** (`fully_tensorized`): Most efficient method that combines networks into a single model ensemble. +4. **Auto-select** (`auto`): Automatically selects the most efficient method based on the number of networks and their architectures. + +### Configuration + +You can specify which training method to use in your experiment configuration: + +```yaml +extra: + training_method: "fully_tensorized" # Options: "auto", "sequential", "tensorized", "fully_tensorized" +``` + +By default, the system will use `"auto"` which automatically selects the most efficient method. + +### Example Config + +An example configuration file is provided at `configs/training_example.yaml` showing how to use the tensorized training methods. + +### Performance + +Our benchmark tests show that the fully tensorized approach can be up to 3x faster than sequential training for larger numbers of networks (e.g., 10+ networks). The speed improvement comes from: + +- Reduced overhead from parallel computation +- Better utilization of GPU resources +- Minimized data transfer between CPU and GPU + +## Tensorized Progressive Dropout + +This document describes the improvements made to the progressive dropout implementation in the alignment codebase, along with benchmark results comparing different approaches. + +### Overview + +Progressive dropout is a technique used to analyze the alignment of neural networks by gradually dropping out neurons and measuring the impact on performance. Three major optimizations are now available: + +1. **Tensorized Implementation**: Process all networks at once using tensor operations +2. **Multi-Strategy Implementation**: Process all strategies (high_rq, low_rq, random) simultaneously +3. **Combined Optimization**: Apply both optimizations together for maximum performance + +### Multi-Strategy Dropout + +The multi-strategy optimization allows processing all three pruning strategies (high_rq, low_rq, random) together, which provides a significant speedup when you need results for all strategies: + +#### Configuration + +```yaml +extra: + use_multi_strategy_dropout: true # Process all strategies simultaneously +``` + +#### Implementation Details + +The implementation creates separate network copies for each strategy, but computes the neuron scores only once. Then it applies different pruning strategies to each network copy and evaluates them in parallel. + +#### Benchmark Results + +Benchmark results for multi-strategy dropout show significant speedups (typically 2.5-3x) compared to running the three strategies sequentially. + +| Configuration | Sequential Time | Multi-Strategy Time | Speedup | +|---------------|----------------|---------------------|---------| +| 5 networks, 10 dropout steps | 240s | 89s | 2.7x | + +### Benchmark Scripts + +Two benchmark scripts are provided to test the performance of these optimizations: + +1. `benchmark_network_training.py`: Compare training methods (sequential, tensorized, fully_tensorized) +2. `benchmark_dropout_strategies.py`: Compare dropout strategy processing (sequential vs. multi-strategy) + +To run the benchmarks: + +```bash +python benchmark_network_training.py --num_networks 10 --epochs 1 +python benchmark_dropout_strategies.py --config configs/config_alignment_experiment.yaml --runs 1 +``` + +### Conclusion + +These optimizations provide substantial performance improvements for alignment experiments, allowing you to run more experiments with larger networks in less time. For the best performance: + +1. Use `training_method: "fully_tensorized"` for training multiple networks +2. Use `use_multi_strategy_dropout: true` when running progressive dropout with multiple strategies \ No newline at end of file diff --git a/_arxiv/_archive/md_files/REFACTORING_METRICS.md b/_arxiv/_archive/md_files/REFACTORING_METRICS.md new file mode 100644 index 00000000..daa2d165 --- /dev/null +++ b/_arxiv/_archive/md_files/REFACTORING_METRICS.md @@ -0,0 +1,66 @@ +# Metrics System Refactoring + +## Overview + +This document summarizes the refactoring of the alignment metrics system, which has been fully consolidated into a single, unified metrics infrastructure under `src/alignment/metrics.py`. + +## Completed Refactoring Tasks + +- [x] Critical Bug Fixes & Import Issues +- [x] Import Hoisting (general pass) +- [x] Logging in DDP Runner +- [x] Optional transformers +- [x] Plotting for Metric Evolution & Conditional Legends +- [x] Configuration Access in ExperimentRunner +- [x] Refactor Pruning Logic in dropout_manager.py +- [x] Consolidated Metrics Systems (metrics_utils.py vs. metrics.py) + - [x] Ported WeightSimilarityMetric to weight_cosine_similarity, weight_dot_similarity, and weight_euclidean_distance + - [x] Ported MIMetric to mi_proj_vs_mean_input + - [x] Ported RQMetric alternative formulation to rq_alt_denom + - [x] Ported NodeRedundancyMetric to node_redundancy + - [x] Updated tests to use the new metrics.py system (get_metric) + - [x] Deprecated and removed metrics_utils.py + +## Architecture Overview + +The new metrics system is built around: + +1. **Central Registry**: `ALIGNMENT_METRICS_REGISTRY` maps metric names to metric functions +2. **Metric Protocol**: `AlignmentMetric` protocol defines the interface for all metrics +3. **Metric Implementation**: `_AlignmentMetricImpl` provides consistent dispatch for all metrics +4. **Factory Function**: `get_metric()` creates properly configured metric objects +5. **High-Level Functions**: + - `compute_metrics_for_layers()`: Process multiple layers at once + - `compute_all_node_scores()`: Process the whole network + - `compute_pairwise_metric()`: Compute metrics between pairs of data + +## Documentation + +A comprehensive metrics system documentation has been added at `src/alignment/README_metrics.md`, which includes: + +- Available metrics and their purposes +- Usage examples for each metric type +- Architecture overview +- Guidelines for metric selection + +## Testing + +All metrics have corresponding tests in `tests/test_benchmark.py` which verify the correct functionality of: + +- Standard RQ and alternative RQ metrics +- MI metrics including the specialized MI projection metric +- Weight similarity metrics +- Node redundancy metrics + +## Future Improvements + +Potential areas for future enhancement: + +1. Add more test coverage for the remaining metrics +2. Consider adding more metrics for specific use cases +3. Performance optimizations for large-scale networks +4. Integration with visualization tools for better insights + +## Conclusion + +The metrics system refactoring has successfully unified all alignment metrics under a single, consistent API. This will make the codebase more maintainable, easier to extend, and provide clearer documentation for users. \ No newline at end of file diff --git a/_arxiv/archive/benchmarks/README.md b/_arxiv/archive/benchmarks/README.md new file mode 100644 index 00000000..5e7498d0 --- /dev/null +++ b/_arxiv/archive/benchmarks/README.md @@ -0,0 +1,29 @@ +# Alignment Benchmarks + +This directory contains benchmarking scripts for the Alignment library. These scripts are designed to assess the performance of various components of the codebase. + +## Contents + +- **benchmark_dropout_strategies.py**: Compares the performance of different dropout strategies (sequential vs. multi-strategy). +- **benchmark_network_training.py**: Benchmarks network training with different configurations. + +## Usage + +Most benchmark scripts can be run directly from the command line: + +```bash +python benchmark_dropout_strategies.py --config configs/config_alignment_experiment.yaml +``` + +## Adding New Benchmarks + +When adding new benchmarks, please follow these guidelines: + +1. Include clear documentation within the script about what is being measured +2. Add command-line arguments for configuration +3. Include a simple way to output/visualize the benchmark results +4. Update this README with information about the new benchmark + +## Results + +Benchmark results should be saved to the `results/` directory, not committed directly to the repository unless they represent an important baseline. \ No newline at end of file diff --git a/_arxiv/archive/benchmarks/benchmark_dropout_strategies.py b/_arxiv/archive/benchmarks/benchmark_dropout_strategies.py new file mode 100755 index 00000000..392afece --- /dev/null +++ b/_arxiv/archive/benchmarks/benchmark_dropout_strategies.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python +""" +Benchmark script to compare original sequential dropout approach vs. multi-strategy approach. + +This script measures the performance difference between: +1. Processing strategies (high_rq, low_rq, random) sequentially +2. Processing all strategies simultaneously with the new tensorized multi-strategy approach +""" + +import os +import sys +import time +import argparse +import logging +import torch +import numpy as np +from tqdm import tqdm + +from alignment.experiments.alignment_experiments import AlignmentExperiment +from alignment.config import ExperimentConfig +from alignment.datasets import load_dataset +from alignment.dropout import progressive_dropout, progressive_dropout_multi_strategy + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s [%(levelname)s] %(message)s', + handlers=[ + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) + +def run_benchmark(config_path, num_runs=3): + """Run benchmark comparing sequential vs. multi-strategy dropout.""" + # Load configuration + logger.info(f"Loading configuration from {config_path}") + config = ExperimentConfig.load(config_path) + + # Force the device to be cuda if available + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + config.device = str(device) + + # Create experiment + experiment = AlignmentExperiment(config) + + # Create networks + logger.info("Creating networks...") + networks = experiment.create_networks() + logger.info(f"Created {len(networks)} networks") + + # Load dataset + logger.info("Loading dataset...") + dataset = load_dataset(config.dataset) + + # Setup dropout parameters + dropout_fractions = np.linspace( + config.alignment.dropout_min, + config.alignment.dropout_max, + config.alignment.dropout_steps + ).tolist() + + pruning_mode = getattr(config.extra, "dropout_pruning_mode", "global_joint") + dropout_mode = getattr(config.extra, "dropout_mode", "scaled") + metric = experiment.metric + + # Initialize timing results + sequential_times = [] + multi_strategy_times = [] + + # Run benchmarks multiple times + for run in range(num_runs): + logger.info(f"\nRun {run+1}/{num_runs}") + + # Make copies of networks to ensure fair comparison + networks_sequential = [net.clone() if hasattr(net, 'clone') else net for net in networks] + networks_multi = [net.clone() if hasattr(net, 'clone') else net for net in networks] + + # Process strategies sequentially (original approach) + logger.info("Running sequential approach...") + sequential_start = time.time() + + strategies = ["high_rq", "low_rq", "random"] + for strategy in tqdm(strategies, desc="Strategies"): + # Clone networks for this strategy to avoid interference + strategy_networks = [net.clone() if hasattr(net, 'clone') else net for net in networks_sequential] + + # Run progressive dropout with this strategy + network_accuracies, network_losses = progressive_dropout( + strategy_networks, + dataset, + dropout_fractions, + metric, + device, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + strategy=strategy, + show_progress=False # Disable progress bars for cleaner output + ) + + sequential_time = time.time() - sequential_start + sequential_times.append(sequential_time) + logger.info(f"Sequential approach took {sequential_time:.2f} seconds") + + # Process all strategies at once (new approach) + logger.info("Running multi-strategy approach...") + multi_start = time.time() + + # Run with multi-strategy mode + network_accuracies, network_losses = progressive_dropout( + networks_multi, + dataset, + dropout_fractions, + metric, + device, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + show_progress=False, # Disable progress bars for cleaner output + use_multi_strategy=True + ) + + multi_time = time.time() - multi_start + multi_strategy_times.append(multi_time) + logger.info(f"Multi-strategy approach took {multi_time:.2f} seconds") + + # Calculate speedup + speedup = sequential_time / multi_time if multi_time > 0 else float('inf') + logger.info(f"Speedup: {speedup:.2f}x") + + # Calculate average times and speedup + avg_sequential = np.mean(sequential_times) + avg_multi = np.mean(multi_strategy_times) + avg_speedup = avg_sequential / avg_multi if avg_multi > 0 else float('inf') + + std_sequential = np.std(sequential_times) + std_multi = np.std(multi_strategy_times) + + # Print summary + logger.info("\n" + "="*50) + logger.info("BENCHMARK SUMMARY") + logger.info("="*50) + logger.info(f"Network count: {len(networks)}") + logger.info(f"Dropout steps: {len(dropout_fractions)}") + logger.info(f"Pruning mode: {pruning_mode}") + logger.info(f"Dropout mode: {dropout_mode}") + logger.info(f"Device: {device}") + logger.info("-"*50) + logger.info(f"Sequential approach: {avg_sequential:.2f} ± {std_sequential:.2f} seconds") + logger.info(f"Multi-strategy approach: {avg_multi:.2f} ± {std_multi:.2f} seconds") + logger.info(f"Average speedup: {avg_speedup:.2f}x") + logger.info("="*50) + + return { + "sequential": sequential_times, + "multi_strategy": multi_strategy_times, + "speedup": avg_speedup + } + +def main(): + parser = argparse.ArgumentParser(description="Benchmark dropout strategies") + parser.add_argument("--config", type=str, default="configs/config_alignment_experiment.yaml", + help="Path to config file") + parser.add_argument("--runs", type=int, default=1, + help="Number of benchmark runs") + + args = parser.parse_args() + + # Run the benchmark + results = run_benchmark(args.config, args.runs) + + # Print the recommendation + if results["speedup"] > 1.5: + print("\nRECOMMENDATION: Use multi-strategy approach for significant speedup.") + else: + print("\nRECOMMENDATION: Both approaches have similar performance in this configuration.") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/_arxiv/archive/benchmarks/benchmark_network_training.py b/_arxiv/archive/benchmarks/benchmark_network_training.py new file mode 100755 index 00000000..72d1b49c --- /dev/null +++ b/_arxiv/archive/benchmarks/benchmark_network_training.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python +""" +Benchmark script for comparing network training methods. + +This script benchmarks different training methods for multiple networks: +1. Sequential training (original approach) +2. Tensorized training (improved parallel approach) +3. Fully tensorized training (optimized ensemble approach) +""" + +import os +import time +import argparse +import logging +import torch +import torch.nn as nn +from tqdm import tqdm + +from alignment.metrics import get_metric +from alignment.datasets import load_dataset +from alignment.models.models import MLP, create_mlp +from alignment.models.base import AlignmentNetwork +from alignment.training import ( + train_networks_sequential, + train_networks_tensorized, + train_networks_fully_tensorized, + train_networks +) + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s [%(levelname)s] %(message)s', + handlers=[ + logging.StreamHandler() + ] +) +logger = logging.getLogger(__name__) + +def create_networks(num_networks, input_size=784, hidden_sizes=[512, 256], output_size=10, seed=42): + """Create multiple networks with identical architecture but different initializations.""" + networks = [] + + for i in range(num_networks): + # Set seed for reproducibility, but different for each network + torch.manual_seed(seed + i) + + # Create a simple MLP network + base_model = MLP( + input_dim=input_size, + output_dim=output_size, + num_hidden=hidden_sizes, + dropout_rate=0.0 + ) + + # Get all linear layers for alignment + linear_layers = {} + for name, module in base_model.named_modules(): + if isinstance(module, nn.Linear): + if name != "layers.0": # Skip the first layer + linear_layers[name] = None # Use the layer's own input + + # Create AlignmentNetwork wrapper + network = AlignmentNetwork(base_model=base_model, alignment_layer_names=linear_layers) + + # Move to device (will be moved again in training, but good practice) + if torch.cuda.is_available(): + network.to('cuda') + + networks.append(network) + + logger.info(f"Created {num_networks} networks with architecture: {input_size} -> {hidden_sizes} -> {output_size}") + return networks + +def run_benchmark(networks, dataset, num_epochs=2, learning_rate=0.001, device=None, show_progress=True): + """Run benchmark comparing different training methods.""" + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + logger.info(f"Running benchmark on device: {device}") + + # Define methods to benchmark + methods = [ + ("Sequential", train_networks_sequential), + ("Tensorized", train_networks_tensorized), + ("Fully Tensorized", train_networks_fully_tensorized), + ("Auto-select", train_networks) + ] + + results = {} + + # Clone networks for each method to ensure fair comparison + for method_name, method_func in methods: + logger.info(f"Benchmarking {method_name} training...") + + # Create copies of networks for this method + method_networks = [] + + # For each network in the original list + for i, source_net in enumerate(networks): + # Get the architecture of the source network + input_dim = 784 # MNIST standard + output_dim = 10 # MNIST standard + + # Extract hidden sizes by examining the linear layers + linear_layers = [ + module for module in source_net.base_model.modules() + if isinstance(module, nn.Linear) + ] + + # Get all but the last layer's output dimensions for hidden sizes + hidden_sizes = [layer.out_features for layer in linear_layers[:-1]] + + logger.debug(f"Network {i}: input_dim={input_dim}, hidden_sizes={hidden_sizes}, output_dim={output_dim}") + + # Create base model with identical architecture + base_model = MLP( + input_dim=input_dim, + output_dim=output_dim, + num_hidden=hidden_sizes, + dropout_rate=0.0 + ) + + # Create alignment layer mapping identical to source network + # Use the layer_to_input_names attribute which is the internal representation + layer_to_input_names = source_net.layer_to_input_names + + # Create new network with same architecture and alignment layers + target_net = AlignmentNetwork( + base_model=base_model, + alignment_layer_names=layer_to_input_names + ) + + # Copy parameters from source to target + # This would be the most reliable way to ensure identical networks + target_net.load_state_dict(source_net.state_dict()) + + # Move to device + target_net.to(device) + + # Add to method networks + method_networks.append(target_net) + + # Measure time + start_time = time.time() + + # Train networks using this method + training_history = method_func( + networks=method_networks, + dataset=dataset, + num_epochs=num_epochs, + learning_rate=learning_rate, + device=device, + show_progress=show_progress + ) + + # Calculate elapsed time + elapsed_time = time.time() - start_time + + # Store results + results[method_name] = { + "time": elapsed_time, + "history": training_history, + "final_acc": training_history["test_acc"][-1] if training_history["test_acc"] else 0.0 + } + + logger.info(f"{method_name} training completed in {elapsed_time:.2f} seconds (final acc: {results[method_name]['final_acc']:.2f}%)") + + # Print summary + print("\nBenchmark Results:") + print("==================") + print(f"Networks: {len(networks)}, Epochs: {num_epochs}") + print("------------------") + + # Find the fastest method to calculate speedup + fastest_time = min(results[method]["time"] for method in results) + + for method_name in results: + elapsed_time = results[method_name]["time"] + final_acc = results[method_name]["final_acc"] + speedup = fastest_time / elapsed_time if elapsed_time > 0 else 0 + + print(f"{method_name:16s}: {elapsed_time:.2f}s ({speedup:.2f}x), Acc: {final_acc:.2f}%") + + return results + +def main(): + parser = argparse.ArgumentParser(description="Benchmark network training methods.") + parser.add_argument("--num_networks", type=int, default=5, help="Number of networks to train") + parser.add_argument("--hidden_sizes", type=str, default="512,256", help="Hidden layer sizes (comma-separated)") + parser.add_argument("--epochs", type=int, default=2, help="Number of epochs to train") + parser.add_argument("--device", type=str, default="cuda", help="Device to run benchmark on") + parser.add_argument("--batch_size", type=int, default=128, help="Batch size for training") + parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility") + + args = parser.parse_args() + + # Parse hidden sizes + hidden_sizes = [int(size) for size in args.hidden_sizes.split(",")] + + # Set device + device = torch.device(args.device if torch.cuda.is_available() and args.device == "cuda" else "cpu") + + # Set random seed + torch.manual_seed(args.seed) + + # Create networks + networks = create_networks( + num_networks=args.num_networks, + hidden_sizes=hidden_sizes, + seed=args.seed + ) + + # Load dataset + logger.info("Loading MNIST dataset") + transform_params = { + "center_crop": None, + "resize": None, + "flatten": True, # Flatten for MLP + "normalize": False + } + dataset_config = { + "dataset_name": "MNIST", + "batch_size": args.batch_size, + "data_path": "./data", + "transform_params": transform_params + } + dataset = load_dataset(dataset_config) + + # Run benchmark + run_benchmark( + networks=networks, + dataset=dataset, + num_epochs=args.epochs, + device=device + ) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/_arxiv/archive/cluster/alignment-ddp-example.slurm b/_arxiv/archive/cluster/alignment-ddp-example.slurm new file mode 100644 index 00000000..13f6226c --- /dev/null +++ b/_arxiv/archive/cluster/alignment-ddp-example.slurm @@ -0,0 +1,157 @@ +#!/bin/bash +#SBATCH --job-name=alignment-ddp # Job name +#SBATCH --partition=kempner # Partition with GPUs +#SBATCH --nodes=2 # Number of nodes +#SBATCH --ntasks-per-node=4 # Tasks per node (match GPUs) +#SBATCH --cpus-per-task=16 # CPUs per task +#SBATCH --mem=512G # Memory per node +#SBATCH --gres=gpu:4 # GPUs per node +#SBATCH --time=24:00:00 # Time limit +#SBATCH --mail-type=begin,end # Email notifications +#SBATCH --account=kempner_lab # Account to charge + +# Set up job directory +JOB_NAME="alignment-experiment" +source cluster/slurm_settings.txt +JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} +mkdir -p $JOB_DIR + +# Set up logging +logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out" +echo "Job started at $(date)" > $logfile +echo "Running on nodes: $SLURM_JOB_NODELIST" >> $logfile + +# Set up DDP environment variables +export MASTER_PORT=12355 +export WORLD_SIZE=$(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE)) +master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_ADDR=$master_addr + +echo "WORLD_SIZE=$WORLD_SIZE" >> $logfile +echo "MASTER_ADDR=$MASTER_ADDR" >> $logfile +echo "MASTER_PORT=$MASTER_PORT" >> $logfile + +# Load modules and activate environment +module purge +module load python/3.9.12 +module load cuda/11.7 + +# Activate conda environment +conda activate alignment_env + +# Record start time +start_time=$(date +%s) + +# Create Python script for the experiment +cat > ${JOB_DIR}/run_experiment.py << 'EOF' +import os +import torch +import torch.distributed as dist +from pathlib import Path + +# Import from alignment codebase +from alignment import ModelWrapper, DatasetWrapper +from alignment.experiments import ProgressiveDropoutExperiment +from alignment.experiments.base import ExperimentConfig +from alignment.infrastructure.computing.distributed import setup_distributed, cleanup_distributed +from alignment.infrastructure.storage.logging import setup_logging + +def main(): + # Setup distributed environment + setup_distributed() + + # Get rank and world size + rank = dist.get_rank() if dist.is_initialized() else 0 + world_size = dist.get_world_size() if dist.is_initialized() else 1 + + # Setup logging (only on main process) + if rank == 0: + setup_logging(Path(os.environ.get('JOB_DIR', '.'))) + + # Configuration + config = ExperimentConfig( + name=f"resnet50_imagenet_ddp", + description="Progressive dropout analysis on ResNet50", + + # Model settings + model_name="resnet50", + pretrained=True, + + # Dataset settings + dataset_name="imagenet", + data_path="/n/holylabs/LABS/kempner_shared/datasets/imagenet", + batch_size=256 // world_size, # Scale batch size per GPU + num_workers=8, + + # Experiment settings + metrics=["rayleigh_quotient", "mi_gaussian", "pid_shared"], + tracked_layers=[ + "layer1.0.conv1", "layer1.1.conv1", "layer1.2.conv1", + "layer2.0.conv1", "layer2.1.conv1", "layer2.2.conv1", + "layer3.0.conv1", "layer3.1.conv1", "layer3.2.conv1", + "layer4.0.conv1", "layer4.1.conv1", "layer4.2.conv1", + ], + + # Training settings + train_before_dropout=True, + training_epochs=90, + learning_rate=0.1 * (world_size * 256 / 256), # Linear LR scaling + optimizer="sgd", + + # Distributed settings + distributed=True, + world_size=world_size, + rank=rank, + + # Other settings + checkpoint_dir=f"{os.environ.get('JOB_DIR', '.')}/checkpoints", + log_dir=f"{os.environ.get('JOB_DIR', '.')}/logs", + exclude_classification_layer=True, + scale_by_norm=False, + force_cpu_for_large_metric_ops=True, + ) + + # Create experiment + experiment = ProgressiveDropoutExperiment( + config=config, + dropout_range=(0.0, 0.9, 10), # 10 dropout levels from 0% to 90% + dropout_mode="magnitude", # Drop based on metric magnitude + compute_metrics_during_training=True, + metric_computation_interval=10, # Compute metrics every 10 epochs + ) + + # Run experiment + print(f"Rank {rank}: Starting experiment...") + results = experiment.run() + + # Save results (only on main process) + if rank == 0: + experiment.save_results() + print("Experiment completed successfully!") + + # Cleanup + cleanup_distributed() + +if __name__ == "__main__": + main() +EOF + +# Run the experiment using srun +# Each process will be launched on the appropriate node/GPU +echo "Starting distributed training..." >> $logfile +srun python ${JOB_DIR}/run_experiment.py >> $logfile 2>&1 + +# Record end time +end_time=$(date +%s) +total_time=$((end_time - start_time)) + +echo "Job completed at $(date)" >> $logfile +echo "Total runtime: $total_time seconds" >> $logfile + +# Copy results to a permanent location (only from rank 0 node) +if [[ $SLURM_PROCID -eq 0 ]]; then + RESULTS_DIR="/n/holylabs/LABS/kempner_lab/Users/$USER/alignment_results/${SLURM_JOB_ID}" + mkdir -p $RESULTS_DIR + cp -r ${JOB_DIR}/* $RESULTS_DIR/ + echo "Results copied to: $RESULTS_DIR" >> $logfile +fi \ No newline at end of file diff --git a/_arxiv/archive/cluster/alignment_stats.slurm b/_arxiv/archive/cluster/alignment_stats.slurm new file mode 100644 index 00000000..5b3d2ad4 --- /dev/null +++ b/_arxiv/archive/cluster/alignment_stats.slurm @@ -0,0 +1,42 @@ +#!/bin/bash +#SBATCH --job-name=alignment_stats # create a short name for your job +#SBATCH --partition=kempner # partition +#SBATCH --account=kempner_bsabatini_lab # account needed for kempner partition +#SBATCH --nodes=1 # node count +#SBATCH --ntasks-per-node=1 # total number of tasks per node +#SBATCH --cpus-per-task=64 # cpu-cores per task (>1 if multi-threaded tasks) +#SBATCH --gres=gpu:1 # number of allocated gpus per node +#SBATCH --mem=1000G # total memory per node (4 GB per cpu-core is default) +#SBATCH --time=04:00:00 # total run time limit (HH:MM:SS) +#SBATCH --mail-type=begin # send email when job begins +#SBATCH --mail-type=end # send email when job ends + +# we need to define the job name directly since it isn't a slurm environment variable +JOB_NAME="alignment_stats" + +# this text file has some settings in it (like the standard job directory) +source cluster/slurm_settings.txt +JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} # make a specific directory for this particular job +mkdir -p $JOB_DIR + +# define a unique log file in the right place +logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out" +echo "Writing to ${logfile}" + +# load python and activate our conda environment +module purge +module load python +conda activate networkAlignmentAnalysis + +# record the start time +start_time=$(date +%s) + +# this is the command that initiates the processes +python experiment.py alignment_stats --network CNN2P2 --dataset CIFAR100 --use_wandb --epochs 200 >> $logfile + +# record the end time +end_time=$(date +%s) + +# measure the time elapsed for the core part of the job in the logfile +total_time=$((end_time-start_time)) +echo "Total Time= "$total_time" seconds" >> $logfile diff --git a/_arxiv/archive/cluster/ddp-example/README.md b/_arxiv/archive/cluster/ddp-example/README.md new file mode 100644 index 00000000..f559b6d8 --- /dev/null +++ b/_arxiv/archive/cluster/ddp-example/README.md @@ -0,0 +1,13 @@ +## Distributed Data Parallel Example on Slurm Cluster + +The contents of this folder contain a MWE for getting DDP to work on a cluster. If you want to try +it yourself, start from the main networkAlignmentAnalysis folder (with the .gitignore etc.). The +file called [ddp.slurm](ddp.slurm) is an ``sbatch`` script you can call with: + +```shell +sbatch cluster/ddp-example/ddp.slurm +``` + +This prepares some environment variables, creates a job folder, then calls the python script +[ddp.py](ddp.py) which downloads MNIST to the job folder, trains it across N GPUs (where N is +equal to the number of nodes times the number of GPUs per node), and saves it if requested. diff --git a/_arxiv/archive/cluster/ddp-example/ddp.py b/_arxiv/archive/cluster/ddp-example/ddp.py new file mode 100644 index 00000000..9eb76ebb --- /dev/null +++ b/_arxiv/archive/cluster/ddp-example/ddp.py @@ -0,0 +1,218 @@ +""" +Massive thank you to the Princeton cluster engineers for providing this example (we've only made small changes) +https://github.com/PrincetonUniversity/multi_gpu_training/blob/main/02_pytorch_ddp/mnist_classify_ddp.py +""" + +import argparse +import shutil +import torch +from torch import nn +import torch.nn.functional as F +import torch.optim as optim +from torchvision import datasets, transforms +from torch.optim.lr_scheduler import StepLR + +import os +import time +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from socket import gethostname + + +class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.conv1 = nn.Conv2d(1, 32, 3, 1) + self.conv2 = nn.Conv2d(32, 64, 3, 1) + self.dropout1 = nn.Dropout(0.25) + self.dropout2 = nn.Dropout(0.5) + self.fc1 = nn.Linear(9216, 128) + self.fc2 = nn.Linear(128, 10) + + def forward(self, x): + x = self.conv1(x) + x = F.relu(x) + x = self.conv2(x) + x = F.relu(x) + x = F.max_pool2d(x, 2) + x = self.dropout1(x) + x = torch.flatten(x, 1) + x = self.fc1(x) + x = F.relu(x) + x = self.dropout2(x) + x = self.fc2(x) + output = F.log_softmax(x, dim=1) + return output + + +def train(args, model, device, dataloader, datasampler, optimizer, epoch, rank): + """basic training script""" + # required for different shuffle order of examples in dataset each epoch + datasampler.set_epoch(epoch) + + if rank == 0 and epoch == 1: + first_batch_timer = time.time() + + model.train() + for batch_idx, (data, target) in enumerate(dataloader): + data, target = data.to(device), target.to(device) + if rank == 0 and epoch == 1 and batch_idx == 0: + print(f"Train-- epoch {epoch}, rank {rank}, first batch loaded in {time.time() - first_batch_timer} seconds.") + optimizer.zero_grad() + output = model(data) + loss = F.nll_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % args.log_interval == 0: + if rank == 0: + print(f"Train Epoch: {epoch} [{batch_idx}/{len(dataloader)} ({100.*batch_idx/len(dataloader):.0f}%)] \t Loss: {loss.item():.6f}") + if args.dry_run: + break + + +def test(model, device, dataloader): + model.eval() + test_loss = 0 + correct = 0 + attempts = 0 + with torch.no_grad(): + for data, target in dataloader: + data, target = data.to(device), target.to(device) + output = model(data) + test_loss += F.nll_loss(output, target, reduction="sum").item() # sum up batch loss + pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability + correct += pred.eq(target.view_as(pred)).sum().item() + attempts += data.size(0) + + test_loss /= attempts + + print(f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{attempts} ({100.*correct/attempts:.0f}%)\n") + + +def setup(rank, world_size): + # initialize the process group + dist.init_process_group("nccl", rank=rank, world_size=world_size) + + +def main(): + # Training settings + parser = argparse.ArgumentParser(description="PyTorch MNIST Example with DDP") + parser.add_argument( + "--job-folder", + type=str, + default=".", + help="job folder for storing data", + ) + parser.add_argument( + "--batch-size", + type=int, + default=64, + metavar="N", + help="input batch size for training (default: 64)", + ) + parser.add_argument( + "--epochs", + type=int, + default=14, + metavar="N", + help="number of epochs to train (default: 14)", + ) + parser.add_argument("--lr", type=float, default=1.0, metavar="LR", help="learning rate (default: 1.0)") + parser.add_argument( + "--gamma", + type=float, + default=0.9, + metavar="M", + help="Learning rate step gamma (default: 0.9)", + ) + parser.add_argument("--no-cuda", action="store_true", default=False, help="disables CUDA training") + parser.add_argument("--dry-run", action="store_true", default=False, help="quickly check a single pass") + parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)") + parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", + ) + parser.add_argument("--save-model", action="store_true", default=False, help="For Saving the current Model") + args = parser.parse_args() + + print("job folder:", args.job_folder) + data_folder = os.path.join(args.job_folder, "data") + + torch.manual_seed(args.seed) + + world_size = int(os.environ["WORLD_SIZE"]) + rank = int(os.environ["SLURM_PROCID"]) + gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"]) + print("gpus_per_node:", gpus_per_node) + print("device_count:", torch.cuda.device_count()) + + assert gpus_per_node == torch.cuda.device_count() + print( + f"Hello from rank {rank} of {world_size} on {gethostname()} where there are" f" {gpus_per_node} allocated GPUs per node.", + flush=True, + ) + + if world_size > 1: + setup(rank, world_size) + if rank == 0: + print(f"Group initialized? {dist.is_initialized()}", flush=True) + + local_rank = rank - gpus_per_node * (rank // gpus_per_node) + torch.cuda.set_device(local_rank) + print(f"host: {gethostname()}, rank: {rank}, local_rank: {local_rank}") + + # Create network + net = Net() + model = net.to(local_rank) + + # Make it a DDP object for distributed processing and training + ddp_model = DDP(model, device_ids=[local_rank]) if world_size > 1 else model + optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr) + scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) + + transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) + + train_data = datasets.MNIST(data_folder, train=True, download=True, transform=transform) + test_data = datasets.MNIST(data_folder, train=False, download=True, transform=transform) + + train_sampler = torch.utils.data.distributed.DistributedSampler(train_data, num_replicas=world_size, rank=rank) + train_loader = torch.utils.data.DataLoader( + train_data, + batch_size=args.batch_size, + sampler=train_sampler, + num_workers=int(os.environ["SLURM_CPUS_PER_TASK"]), + pin_memory=True, + ) + test_loader = torch.utils.data.DataLoader( + test_data, + batch_size=args.batch_size, + num_workers=int(os.environ["SLURM_CPUS_PER_TASK"]), + pin_memory=True, + ) + + for epoch in range(1, args.epochs + 1): + if rank == 0: + epoch_time = time.time() + train(args, ddp_model, local_rank, train_loader, train_sampler, optimizer, epoch, rank) + if rank == 0: + test(ddp_model, local_rank, test_loader) + scheduler.step() + if rank == 0: + epoch_time = time.time() - epoch_time + print(f"\nEpoch {epoch}, Train & Test Time = {epoch_time:.1f} seconds (measured from rank {rank}).\n") + + if args.save_model and rank == 0: + torch.save(model.state_dict(), os.path.join(args.job_folder, "test_model_ddp.pt")) + + if world_size > 1: + dist.destroy_process_group() + + # clear locally downloaded MNIST data + shutil.rmtree(data_folder) + + +if __name__ == "__main__": + main() diff --git a/_arxiv/archive/cluster/ddp-example/ddp.slurm b/_arxiv/archive/cluster/ddp-example/ddp.slurm new file mode 100644 index 00000000..43cdd746 --- /dev/null +++ b/_arxiv/archive/cluster/ddp-example/ddp.slurm @@ -0,0 +1,56 @@ +#!/bin/bash +#SBATCH --job-name=ddp-example # create a short name for your job +#SBATCH --partition=kempner # partition to use (need one with GPUs for this) +#SBATCH --nodes=2 # node count +#SBATCH --ntasks-per-node=4 # total number of tasks per node +#SBATCH --cpus-per-task=16 # cpu-cores per task (>1 if multi-threaded tasks) +#SBATCH --mem=512G # total memory per node (4 GB per cpu-core is default) +#SBATCH --gres=gpu:4 # number of allocated gpus per node +#SBATCH --time=00:10:00 # total run time limit (HH:MM:SS) +#SBATCH --mail-type=begin # send email when job begins +#SBATCH --mail-type=end # send email when job ends +#SBATCH --account=kempner_bsabatini_lab # the account needs to be specified for the kempner partition + +# we need to define the job name directly since it isn't a slurm environment variable +JOB_NAME="DDP-Example" + +# this text file has some settings in it (like the standard job directory) +source cluster/slurm_settings.txt +JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} # make a specific directory for this particular job +mkdir -p $JOB_DIR + +# define a unique log file in the right place +logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out" +echo $logfile + +# note: +# it's important for ntasks-per-node to be equal to the number of GPUs per node, which is set with --gres=gpu:4 +export MASTER_PORT=12355 # there may be a smarter way to set this (e.g. with code), but this port is almost always open (maybe 100% of the time) +export WORLD_SIZE=$(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE)) # world size is equal to number of nodes and number of tasks per node +echo "WORLD_SIZE="$WORLD_SIZE >> $logfile +echo "MASTER_PORT="$MASTER_PORT >> $logfile + +# define a master address for communication between GPUs +master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_ADDR=$master_addr +echo "MASTER_ADDR="$MASTER_ADDR >> $logfile + +# load python and activate our conda environment +module purge +module load python +conda activate networkAlignmentAnalysis + +# record the start time +start_time=$(date +%s) + +# this is the command that initiates the processes +# srun will send each command to each task on each node for a total of WORLD_SIZE tasks +# (in this case the script it runs is python cluster/ddp.py) +srun python cluster/ddp-example/ddp.py --epochs=50 --job-folder=${JOB_DIR} >> $logfile + +# record the end time +end_time=$(date +%s) + +# measure the time elapsed for the core part of the job in the logfile +total_time=$((end_time-start_time)) +echo "Total Time= "$total_time" seconds" >> $logfile diff --git a/_arxiv/archive/cluster/imagenet-example/imagenet.py b/_arxiv/archive/cluster/imagenet-example/imagenet.py new file mode 100644 index 00000000..08a2abfd --- /dev/null +++ b/_arxiv/archive/cluster/imagenet-example/imagenet.py @@ -0,0 +1,217 @@ +import argparse +import wandb +import torch +import torch.nn.functional as F +import torch.optim as optim +from torch.optim.lr_scheduler import StepLR + +import os +import sys +import time +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP +from socket import gethostname + +mainPath = os.path.dirname(os.path.abspath(__file__)) + "/.." +sys.path.append(mainPath) + +from networkAlignmentAnalysis import datasets +from networkAlignmentAnalysis.models.registry import get_model + + +def configure_wandb(args): + """create a wandb run file and set environment parameters appropriately""" + if args.use_wandb: + wandb.login() + run = wandb.init( + project=args.job_name, + name=args.job_folder, + ) + os.environ["WANDB_MODE"] = "offline" + return run + + return None + + +def train(args, model, device, dataset, optimizer, epoch, rank, run, train=True): + dataloader = dataset.train_loader if train else dataset.test_loader + if dataset.distributed: + if train: + dataset.train_sampler.set_epoch(epoch) + else: + dataset.test_sampler.set_epoch(epoch) + + if rank == 0: + first_batch_timer = time.time() + + model.train() + for batch_idx, batch in enumerate(dataloader): + data, target = dataset.unwrap_batch(batch, device=device) + if rank == 0 and batch_idx == 0: + print(f"Train-- epoch {epoch}, rank {rank}, first batch loaded in {time.time() - first_batch_timer} seconds.") + optimizer.zero_grad() + output = model(data) + loss = dataset.measure_loss(output, target) + loss.backward() + optimizer.step() + if batch_idx % args.log_interval == 0: + if rank == 0: + print(f"Train Epoch: {epoch} [{batch_idx}/{len(dataloader)} ({100.*batch_idx/len(dataloader):.0f}%)] \t Loss: {loss.item():.6f}") + if run is not None: + run.log(dict(epoch=epoch, batch_idx=batch_idx, train_loss=loss.item())) + if args.dry_run: + break + + +def test(model, device, dataset, run, train=False): + dataloader = dataset.train_loader if train else dataset.test_loader + model.eval() + test_loss = 0 + correct = 0 + attempts = 0 + with torch.no_grad(): + for batch in dataloader: + data, target = dataset.unwrap_batch(batch, device=device) + output = model(data) + test_loss += dataset.measure_loss(output, target, reduction="sum").item() + pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability + correct += pred.eq(target.view_as(pred)).sum().item() + attempts += data.size(0) + + test_loss /= attempts + + print(f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{attempts} ({100.*correct/attempts:.0f}%)\n") + + if run is not None: + run.log(dict(test_loss=test_loss, test_accuracy=100.0 * correct / attempts)) + + +def create_dataset(name, net, distributed=True, loader_parameters={}): + return datasets.get_dataset( + name, + build=True, + distributed=distributed, + transform_parameters=net, + loader_parameters=loader_parameters, + ) + + +def setup(rank, world_size): + # initialize the process group + dist.init_process_group("nccl", rank=rank, world_size=world_size) + + +def main(): + # Training settings + parser = argparse.ArgumentParser(description="PyTorch ImageNet Example") + parser.add_argument( + "--job-name", + type=str, + default="ImageNet-Example", + help="job name for this project", + ) + parser.add_argument( + "--job-folder", + type=str, + default=".", + help="job folder for storing data", + ) + parser.add_argument( + "--use_wandb", + default=False, + action="store_true", + help="if used, will log experiment to WandB", + ) + parser.add_argument( + "--batch-size", + type=int, + default=64, + metavar="N", + help="input batch size for training (default: 64)", + ) + parser.add_argument( + "--epochs", + type=int, + default=14, + metavar="N", + help="number of epochs to train (default: 14)", + ) + parser.add_argument("--lr", type=float, default=1.0, metavar="LR", help="learning rate (default: 1.0)") + parser.add_argument("--wd", type=float, default=0.0, metavar="WD", help="weight decay (default: 0.0)") + parser.add_argument( + "--gamma", + type=float, + default=0.9, + metavar="M", + help="Learning rate step gamma (default: 0.99)", + ) + parser.add_argument("--no-cuda", action="store_true", default=False, help="disables CUDA training") + parser.add_argument("--dry-run", action="store_true", default=False, help="quickly check a single pass") + parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)") + parser.add_argument( + "--log-interval", + type=int, + default=10, + metavar="N", + help="how many batches to wait before logging training status", + ) + parser.add_argument("--save-model", action="store_true", default=False, help="For Saving the current Model") + args = parser.parse_args() + + run = configure_wandb(args) + + torch.manual_seed(args.seed) + + world_size = int(os.environ["WORLD_SIZE"]) + rank = int(os.environ["SLURM_PROCID"]) + gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"]) + assert gpus_per_node == torch.cuda.device_count() + print( + f"Hello from rank {rank} of {world_size} on {gethostname()} where there are" f" {gpus_per_node} allocated GPUs per node.", + flush=True, + ) + + loader_parameters = dict( + batch_size=args.batch_size, + num_workers=max(int(os.environ["SLURM_CPUS_PER_TASK"]) - 1, 1), # save a little headroom + ) + + if world_size > 1: + setup(rank, world_size) + if rank == 0: + print(f"Group initialized? {dist.is_initialized()}", flush=True) + + local_rank = rank - gpus_per_node * (rank // gpus_per_node) + torch.cuda.set_device(local_rank) + print(f"host: {gethostname()}, rank: {rank}, local_rank: {local_rank}") + + model_name = "AlexNet" + dataset_name = "ImageNet" + net = get_model(model_name, build=True, dataset=dataset_name) + dataset = create_dataset(dataset_name, net, distributed=world_size > 1, loader_parameters=loader_parameters) + + model = net.to(local_rank) + ddp_model = DDP(model, device_ids=[local_rank]) if world_size > 1 else model + optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr, weight_decay=args.wd) + + scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) + + for epoch in range(1, args.epochs + 1): + epoch_time = time.time() + train(args, ddp_model, local_rank, dataset, optimizer, epoch, rank, run) + if rank == 0: + test(ddp_model, local_rank, dataset, run) + scheduler.step() + epoch_time = time.time() - epoch_time + if rank == 0: + print(f"Epoch {epoch}, Train & Test Time = {epoch_time:.1f} seconds (measured from rank {rank}).\n") + + if args.save_model and rank == 0: + torch.save(model.state_dict(), os.path.join(args.job_folder, "alexnet_imagenet.pt")) + + if world_size > 1: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/_arxiv/archive/cluster/imagenet-example/imagenet.slurm b/_arxiv/archive/cluster/imagenet-example/imagenet.slurm new file mode 100644 index 00000000..fe4e8976 --- /dev/null +++ b/_arxiv/archive/cluster/imagenet-example/imagenet.slurm @@ -0,0 +1,50 @@ +#!/bin/bash +#SBATCH --job-name=ImageNet-Example # create a short name for your job +#SBATCH --partition=kempner # partition +#SBATCH --account=kempner_bsabatini_lab # account needed for kempner partition +#SBATCH --nodes=4 # node count +#SBATCH --ntasks-per-node=4 # total number of tasks per node +#SBATCH --cpus-per-task=16 # cpu-cores per task (>1 if multi-threaded tasks) +#SBATCH --gres=gpu:4 # number of allocated gpus per node +#SBATCH --mem=1000G # total memory per node (4 GB per cpu-core is default) +#SBATCH --time=01:00:00 # total run time limit (HH:MM:SS) +#SBATCH --mail-type=begin # send email when job begins +#SBATCH --mail-type=end # send email when job ends + + +# we need to define the job name directly since it isn't a slurm environment variable +JOB_NAME="ImageNet-Example" + +# this text file has some settings in it (like the standard job directory) +source cluster/slurm_settings.txt +JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} # make a specific directory for this particular job +mkdir -p $JOB_DIR + +# define a unique log file in the right place +logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out" +echo "Writing to "$logfile + +export MASTER_PORT=12355 +export WORLD_SIZE=$(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE)) +echo "WORLD_SIZE="$WORLD_SIZE >> $logfile +echo "MASTER_PORT="$MASTER_PORT >> $logfile + +master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) +export MASTER_ADDR=$master_addr +echo "MASTER_ADDR="$MASTER_ADDR >> $logfile + +module purge +module load python +conda activate networkAlignmentAnalysis + +start_time=$(date +%s) + +# (nvidia-smi --query-gpu=utilization.gpu --format=csv --loop=1 --filename=gpu_utilization.log) + +srun python cluster/imagenet-example/imagenet.py --epochs=10 --use_wandb --log-interval=100 --job-folder=${JOB_DIR} --job-name=${JOB_NAME} --save-model >> $logfile + +end_time=$(date +%s) + +total_time=$((end_time-start_time)) +echo "Total Time= "$total_time" seconds" >> $logfile +echo "Total Time= "$total_time" seconds" diff --git a/_arxiv/archive/cluster/lightning-example/lightning_train.sbatch b/_arxiv/archive/cluster/lightning-example/lightning_train.sbatch new file mode 100644 index 00000000..86a127f2 --- /dev/null +++ b/_arxiv/archive/cluster/lightning-example/lightning_train.sbatch @@ -0,0 +1,35 @@ +#! /bin/bash +#SBATCH --job-name=vast_benchmark +#SBATCH --time=1-00:00:00 +#SBATCH --partition=kempner_requeue +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=4 +#SBATCH --mem=1000G +#SBATCH --output=benchmark_lightning.out +#SBATCH --error=benchmark_lightning.err +#SBATCH --account=kempner_dev +#SBATCH --dependency=singleton +#SBATCH --gres=gpu:1 + +if [[ -z $1 ]] ; then + model="alexnet" +else + model=$1 +fi + + +VAST_CREDENTIALS_PATH=/n/holylabs/LABS/kempner_dev/Users/$USER/vast-benchmarking/.load_credentials.sh +METRIC_GATHERING_SCRIPT_PATH=/n/holylabs/LABS/kempner_dev/Users/$USER/vast-benchmarking/get_vast_stats.sh +DATADIR=/n/fas-vast/kempner_benchmark/preproccessed_imagenet/imagenet21k_resized +CONTAINER_PATH=/n/holylabs/LABS/kempner_dev/Lab/containers/pytorch_2.1.2-cuda12.1-cudnn8-runtime-lightning.sif + + +source $VAST_CREDENTIALS_PATH +metrics_file=vast_metrics/vast_metrics_${SLURM_NNODES}nodes_${SLURM_GPUS_ON_NODE}gpus_${model}_${SLURM_JOB_ID}.csv + +$METRIC_GATHERING_SCRIPT_PATH $metrics_file 30 & + +time srun singularity exec --bind $DATADIR --nv $CONTAINER_PATH python3 train_model_lightning.py --dataset $DATADIR/imagenet21k_train \ + --epochs 1 --batch_size 64 --num_workers $SLURM_CPUS_PER_TASK --num_gpus $SLURM_GPUS_ON_NODE --num_nodes $SLURM_NNODES \ + --model $model \ No newline at end of file diff --git a/_arxiv/archive/cluster/lightning-example/submit_lightning.sh b/_arxiv/archive/cluster/lightning-example/submit_lightning.sh new file mode 100644 index 00000000..3e246f8f --- /dev/null +++ b/_arxiv/archive/cluster/lightning-example/submit_lightning.sh @@ -0,0 +1,26 @@ +#! /bin/bash +# Script used to submit a matrix of jobs to the cluster + +mkdir -p logs +mkdir -p vast_metrics + +for model in alexnet resnet50 ; do + for nodes in 1 2 4 8 16 32 ; do + if [[ $nodes -eq 1 ]]; then + for gpus in 1 2 4 ; do + sbatch -N $nodes --ntasks-per-node $gpus --gres gpu:$gpus \ + -o logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.out \ + -e logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.err \ + lightning_train.sbatch $model + sleep 3 + done + else + gpus=4 + sbatch -N $nodes --ntasks-per-node $gpus --gres gpu:$gpus \ + -o logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.out \ + -e logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.err \ + lightning_train.sbatch $model + sleep 3 + fi + done +done \ No newline at end of file diff --git a/_arxiv/archive/cluster/lightning-example/train_model_lightning.py b/_arxiv/archive/cluster/lightning-example/train_model_lightning.py new file mode 100644 index 00000000..3589f688 --- /dev/null +++ b/_arxiv/archive/cluster/lightning-example/train_model_lightning.py @@ -0,0 +1,110 @@ +from typing import Any +from lightning.pytorch.utilities.types import STEP_OUTPUT +import torch +import torchvision +from torchvision import datasets, transforms +from torch.utils.data import DataLoader + +import lightning as L + +import argparse +import logging + + +class LightningModel(L.LightningModule): + def __init__(self, model, optimizer, lr) -> None: + super().__init__() + self.model = model + self.optimizer = optimizer + self.lr = lr + + def training_step(self, batch, batch_idx): + images, labels = batch + output = self.model(images) + loss = torch.nn.functional.cross_entropy(output, labels) + return loss + + def configure_optimizers(self): + optimizer = self.optimizer(self.parameters(), lr=self.lr) + return optimizer + + +# def train_one_epoch(model, loader, optimizer, device): +# model.train() +# batch_num = 0 +# for images, labels in loader: +# images, labels = images.to(device), labels.to(device) +# optimizer.zero_grad() +# output = model(images) +# loss = torch.nn.functional.cross_entropy(output, labels) +# loss.backward() +# optimizer.step() +# if batch_num % 100 == 0: +# logging.info(f"Batch Num: {batch_num} Loss: {loss.item()}") +# batch_num += 1 + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Train a model") + parser.add_argument("--dataset", type=str, default="imagenet21k_train", help="dataset to use") + parser.add_argument("--batch_size", type=int, default=64, help="batch size") + parser.add_argument("--epochs", type=int, default=10, help="number of epochs") + parser.add_argument("--lr", type=float, default=0.01, help="learning rate") + parser.add_argument("--model", type=str, default="alexnet", help="model to use") + parser.add_argument("--num_workers", type=int, default=4, help="number of workers") + parser.add_argument("--num_gpus", type=int, default=1, help="number of gpus per node") + parser.add_argument("--num_nodes", type=int, default=1, help="number of nodes") + parser.add_argument( + "--checkpoint_steps", + type=int, + default=0, + help="number of training steps to save checkpoint after", + ) + parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="directory to save checkpoints") + + return parser.parse_args() + + +def main(): + args = parse_args() + preprocess = transforms.Compose( + [ + transforms.Resize(256), + transforms.CenterCrop(224), + transforms.ToTensor(), + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ) + dataset = datasets.ImageFolder(args.dataset, transform=preprocess) + logging.info(f"Data loaded, Found {len(dataset.classes)} classes, {len(dataset)} images") + num_classes = len(dataset.classes) + loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) + model = torchvision.models.get_model(args.model, num_classes=num_classes) + optimizer = torch.optim.SGD + lighting_model = LightningModel(model, optimizer, args.lr) + + callbacks = None + if args.checkpoint_steps > 0: + callbacks = [ + L.callbacks.ModelCheckpoint( + dirpath=args.checkpoint_dir, + save_top_k=-1, + save_last=True, + every_n_train_steps=args.checkpoint_steps, + ) + ] + + trainer = L.Trainer( + max_epochs=args.epochs, + min_epochs=args.epochs, + strategy="ddp", + devices=args.num_gpus, + num_nodes=args.num_nodes, + callbacks=callbacks, + ) + trainer.fit(lighting_model, train_dataloaders=loader) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + main() diff --git a/_arxiv/archive/cluster/slurm_settings.txt b/_arxiv/archive/cluster/slurm_settings.txt new file mode 100644 index 00000000..229c3711 --- /dev/null +++ b/_arxiv/archive/cluster/slurm_settings.txt @@ -0,0 +1,3 @@ +# Job Folder +JOB_FOLDER=./jobs + diff --git a/_arxiv/archive/cluster/testing-examples/test_job.slurm b/_arxiv/archive/cluster/testing-examples/test_job.slurm new file mode 100644 index 00000000..7ab81e4d --- /dev/null +++ b/_arxiv/archive/cluster/testing-examples/test_job.slurm @@ -0,0 +1,32 @@ +#!/bin/bash +#SBATCH --job-name=test_files # create a short name for your job +#SBATCH --nodes=1 # node count +#SBATCH --ntasks-per-node=1 # total number of tasks per node +#SBATCH --cpus-per-task=1 # cpu-cores per task (>1 if multi-threaded tasks) +#SBATCH --mem=100M # total memory per node (4 GB per cpu-core is default) +#SBATCH --time=00:20:00 # total run time limit (HH:MM:SS) +#SBATCH --mail-type=begin # send email when job begins +#SBATCH --mail-type=end # send email when job ends +#SBATCH --output=/dev/null # no output message + +source cluster/slurm_settings.txt +mkdir -p $JOB_FOLDER +logfile="${JOB_FOLDER}/slurm-${SLURM_JOB_ID}.out" + +echo "Writing to ${logfile}" +echo "Writing in ${logfile}" >> $logfile + +module purge +module load python + +start_time=$(date +%s) + +python -c "print('hello world')" >> $logfile +python test_print.py >> $logfile + +end_time=$(date +%s) + +total_time=$((end_time-start_time)) +echo "Total Time= "$total_time" seconds" >> $logfile + +echo "Finished." diff --git a/_arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md b/_arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md new file mode 100644 index 00000000..58e65f3b --- /dev/null +++ b/_arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md @@ -0,0 +1,84 @@ +# Gaussian Mutual Information with Edgeworth Expansion + +## Overview +I've implemented a new metric `gaussian_mi_analytic` that computes mutual information between inputs and outputs of neural network nodes, assuming approximately Gaussian distributions with analytic expansions for non-Gaussian corrections. + +## Features + +### 1. **Analytic Gaussian MI Calculation** +For linear transformations Y = WX + ε where X and ε are Gaussian: +- Exact formula: I(X;Y) = 1/2 * log(det(Σ_X) * det(Σ_Y) / det(Σ_joint)) +- Efficient computation using covariance matrices +- Numerical stability with regularization + +### 2. **Edgeworth Expansion Corrections** +The metric includes corrections for non-Gaussian distributions up to order 3: + +- **Order 0**: Pure Gaussian assumption +- **Order 1**: First-order correction using third cumulants (skewness) +- **Order 2**: Second-order correction using fourth cumulants (kurtosis) and mixed terms +- **Order 3**: Third-order correction with cross-terms between skewness and kurtosis + +### 3. **Key Parameters** +- `expansion_order`: Controls the order of Edgeworth expansion (0-3) +- `noise_std`: Assumed noise level in the system +- `regularization`: Numerical stability parameter +- `per_neuron`: Compute MI for each neuron separately or jointly + +## Mathematical Foundation + +### Cumulants Used +- κ₂ = variance +- κ₃ = E[X³] (related to skewness) +- κ₄ = E[X⁴] - 3σ⁴ (excess kurtosis) + +### Edgeworth Corrections +The corrections capture deviations from Gaussianity: + +1. **First-order**: Involves normalized third cumulants (γ₁ = κ₃/σ³) +2. **Second-order**: Involves normalized fourth cumulants (γ₂ = κ₄/σ⁴) +3. **Third-order**: Cross-terms between different order cumulants + +## Test Results + +The metric was tested on: +1. **Pure Gaussian data**: Corrections are minimal as expected +2. **Skewed data (Chi-squared)**: Small corrections observed +3. **Heavy-tailed data (Student-t)**: Significant corrections, especially at orders 2 and 3 +4. **Different noise levels**: MI decreases with increasing noise as expected + +### Example Results +- Pure Gaussian (Order 0): MI ≈ 2.99 +- Heavy-tailed (Order 0): MI ≈ 2.99 +- Heavy-tailed (Order 2): MI ≈ 4.16 (showing significant correction) + +## Usage Example + +```python +from src.alignment.core.registry import METRIC_REGISTRY + +# Create metric with second-order corrections +metric = METRIC_REGISTRY.get("gaussian_mi_analytic")( + expansion_order=2, + noise_std=0.1, + per_neuron=True +) + +# Compute MI scores for each neuron +mi_scores = metric.compute(inputs=inputs, weights=weights) +``` + +## Benefits + +1. **Analytic computation**: Fast and exact for Gaussian case +2. **Non-Gaussian handling**: Edgeworth expansion captures deviations +3. **Flexible**: Can adjust expansion order based on data characteristics +4. **Per-neuron analysis**: Can analyze individual neuron information transfer + +## Integration + +The metric is fully integrated into the alignment framework: +- Registered as `gaussian_mi_analytic` +- Available through the metric registry +- Follows the standard BaseMetric interface +- Properly handles different tensor dimensions \ No newline at end of file diff --git a/_arxiv/archive/scripts/README.md b/_arxiv/archive/scripts/README.md new file mode 100644 index 00000000..f5b9f704 --- /dev/null +++ b/_arxiv/archive/scripts/README.md @@ -0,0 +1,35 @@ +# Alignment Scripts + +This directory contains utility scripts for the Alignment library. These scripts provide various functionality for running experiments, demonstrations, and utilities. + +## Contents + +- **direct_pruning_test.py**: Tests pruning functionality directly without relying on the experiment infrastructure. +- **run_multi_strategy_experiment.py**: Runs experiments with multiple pruning strategies. +- **run_fixed_experiment.py**: Runs experiments with fixed configurations. +- **run_cascading_with_plots.py**: Runs cascading pruning experiments and generates plots. +- **run_cascading_test.py**: Tests cascading pruning methodology. +- **continual_mnist.py**: Implementation of continual learning on MNIST. +- **teacher_student.py**: Implementation of teacher-student neural network model. + +## Shell Scripts + +- **run_benchmark.sh**: Convenience script for running benchmarks. +- **run_cascading_test.sh**: Shell script for testing cascading pruning. + +## Usage + +Most scripts have command-line arguments for configuration: + +```bash +python scripts/run_fixed_experiment.py --config configs/config_alignment_experiment.yaml +``` + +## Adding New Scripts + +When adding new scripts, please follow these guidelines: + +1. Include clear documentation at the top of the script about its purpose and usage +2. Add command-line arguments for configuration using argparse +3. Include appropriate error handling and logging +4. Update this README with information about the new script \ No newline at end of file diff --git a/_arxiv/archive/scripts/__init__.py b/_arxiv/archive/scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/_arxiv/archive/scripts/continual_mnist.py b/_arxiv/archive/scripts/continual_mnist.py new file mode 100644 index 00000000..021b5a55 --- /dev/null +++ b/_arxiv/archive/scripts/continual_mnist.py @@ -0,0 +1,144 @@ +# Some Code for continual learning with permuted MNIST +import time +from functools import partial +import numpy as np +import scipy as sp +import sklearn +import torch +import torch.nn.functional as F +from torch import nn +from torchvision.transforms import v2 as transforms + +from tqdm import tqdm +import matplotlib as mpl +from matplotlib import pyplot as plt + +import os +import sys + +mainPath = os.path.dirname(os.path.abspath(__file__)) + "/../.." +sys.path.append(mainPath) + +from alignment.models.registry import get_model +from alignment.datasets import get_dataset +from alignment.experiments.registry import get_experiment +from alignment import utils +from alignment import files +from alignment import train + + +def permute(batch, batch_dim=True, shuffle_idx=None): + if shuffle_idx is not None: + original_size = batch[0].shape + if batch_dim: + batch[0] = batch[0][:, shuffle_idx] + else: + batch[0] = batch[0][shuffle_idx] + batch[0].reshape(original_size) + return batch + + +def add_permutation(dataset, num_pixels=784): + perm = partial(permute, batch_dim=True, shuffle_idx=torch.randperm(num_pixels)) + if dataset.extra_transform is None: + dataset.extra_transform = [] + dataset.extra_transform.append(perm) + return dataset + + +def update_permutation(dataset, num_pixels=784): + perm = partial(permute, batch_dim=True, shuffle_idx=torch.randperm(num_pixels)) + dataset.extra_transform[-1] = perm + return dataset + + +def do_permuted_round(nets, optimizers, dataset, train_epochs=1, verbose=False): + parameters = dict( + verbose=verbose, + num_epochs=train_epochs, + alignment=False, + ) + dataset = update_permutation(dataset) + + # train and test + train_results = train.train(nets, optimizers, dataset, **parameters) + test_results = train.test(nets, dataset, **parameters) + + return train_results, test_results + + +def main(): + DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + print("using device: ", DEVICE) + + model_name = "MLP" + dataset_name = "MNIST" + + hidden_widths = [200, 200, 200] + + lrs = [1e-1, 3e-2, 1e-2, 3e-3] + num_replicates = 3 + + nets = [] + optimizers = [] + net_lr = [] + for lr in lrs: + for _ in range(num_replicates): + net = get_model(model_name, build=True, dataset=dataset_name, hidden_widths=hidden_widths, dropout=0.0, ignore_flag=False) + net.to(DEVICE) + + optimizer = torch.optim.SGD(net.parameters(), lr=lr) + + nets.append(net) + optimizers.append(optimizer) + net_lr.append(lr) + + loader_parameters = dict( + shuffle=True, + batch_size=5, + ) + dataset = get_dataset(dataset_name, build=True, transform_parameters=net, loader_parameters=loader_parameters, device=DEVICE) + dataset = add_permutation(dataset) + + num_rounds = 200 + + train_loss = [] + train_accuracy = [] + test_loss = [] + test_accuracy = [] + for round in tqdm(range(num_rounds)): + c_train_res, c_test_res = do_permuted_round(nets, optimizers, dataset, verbose=False) + train_loss.append(c_train_res["loss"]) + train_accuracy.append(c_train_res["accuracy"]) + test_loss.append(c_test_res["loss"]) + test_accuracy.append(c_test_res["accuracy"]) + + loss = torch.stack([torch.tensor(l) for l in test_loss]) + accuracy = torch.stack([torch.tensor(a) for a in test_accuracy]) + + type_loss = utils.compute_stats_by_type(loss, len(lrs), 1)[0] + type_accuracy = utils.compute_stats_by_type(accuracy, len(lrs), 1)[0] + + # print(loss.shape, accuracy.shape, type_loss.shape, type_accuracy.shape) + + cols = mpl.colormaps["Set1"].resampled(len(lrs)) + names = [f"lr={lr}" for lr in lrs] + + fig, ax = plt.subplots(1, 2, figsize=(6, 3), layout="constrained") + for ii in range(len(lrs)): + ax[0].plot(range(num_rounds), type_loss[:, ii], c=cols(ii), label=names[ii]) + ax[1].plot(range(num_rounds), type_accuracy[:, ii], c=cols(ii), label=names[ii]) + + ax[0].set_xlabel("Rounds of Permuted MNIST") + ax[1].set_xlabel("Rounds of Permuted MNIST") + ax[0].set_ylabel("loss") + ax[1].set_ylabel("accuracy") + ax[0].legend(loc="best") + ax[1].legend(loc="best") + plt.show() + + print("Hello, add a debugger here to evaluate while testing") + + +if __name__ == "__main__": + main() diff --git a/_arxiv/archive/scripts/direct_pruning_test.py b/_arxiv/archive/scripts/direct_pruning_test.py new file mode 100644 index 00000000..efb97975 --- /dev/null +++ b/_arxiv/archive/scripts/direct_pruning_test.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python +""" +Direct Pruning Test + +This script tests pruning functionality directly, without relying on the complex +experiment infrastructure. It loads a model, applies pruning, and logs the results +including before/after weights and accuracies. +""" + +import os +import sys +import logging +import argparse +import copy +import numpy as np +import torch +import torch.nn as nn + +# Set up logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s [%(levelname)s] %(message)s', + handlers=[ + logging.StreamHandler(), + logging.FileHandler('direct_pruning_test.log', mode='w') + ] +) +logger = logging.getLogger(__name__) + +# Add src to path if not already there +src_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src') +if src_path not in sys.path: + sys.path.insert(0, src_path) + +# Import alignment modules +from alignment.models.registry import create_model +from alignment.datasets import load_dataset +from alignment.metrics import get_metric +from alignment.dropout import progressive_dropout +from alignment.config import ExperimentConfig + +def count_zero_weights(model): + """Count zero weights in the model.""" + zero_weights = 0 + total_weights = 0 + + for name, param in model.named_parameters(): + if 'weight' in name: + layer_zeros = (param.data == 0).sum().item() + layer_total = param.data.numel() + zero_weights += layer_zeros + total_weights += layer_total + + if hasattr(param, 'shape'): + logger.info(f"Layer {name}: {layer_zeros}/{layer_total} zeros " + f"({100.0*layer_zeros/layer_total:.2f}% pruned), shape {param.shape}") + + if total_weights > 0: + logger.info(f"Total: {zero_weights}/{total_weights} zeros " + f"({100.0*zero_weights/total_weights:.2f}% pruned)") + + return zero_weights, total_weights + +def verify_pruning(config_path=None, model_name='mlp', dataset_name='mnist', strategy='low_rq', + pruning_mode='layer_wise', dropout_mode='unscaled', pruning_percent=0.5, + device=None): + """Test pruning directly.""" + # Use GPU if available + if device is None: + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + logger.info(f"Using device: {device}") + + # Fix dataset name casing - use uppercase for standard datasets + if dataset_name.lower() == 'mnist': + dataset_name = 'MNIST' + elif dataset_name.lower() == 'cifar10': + dataset_name = 'CIFAR10' + elif dataset_name.lower() == 'cifar100': + dataset_name = 'CIFAR100' + + # Define data path + data_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data') + + # Create dataset config + dataset_config = { + 'dataset_name': dataset_name, + 'batch_size': 128, + 'data_dir': data_root, + 'download': True + } + + # Make sure data directory exists + os.makedirs(data_root, exist_ok=True) + + dataset = load_dataset(dataset_config) + + # Get input dimension from dataset + if dataset_name.lower() in ['mnist', 'fashion_mnist']: + input_dim = 784 # 28*28 + else: + # Try to infer from the dataset + sample_data = next(iter(dataset.train_loader))[0] + input_dim = sample_data[0].flatten().shape[0] + logger.info(f"Inferred input dimension: {input_dim}") + + # Create a simple model + class MLP(nn.Module): + def __init__(self, input_dim=784, hidden_dim=100, output_dim=10): + super().__init__() + self.fc1 = nn.Linear(input_dim, hidden_dim) + self.fc2 = nn.Linear(hidden_dim, hidden_dim) + self.fc3 = nn.Linear(hidden_dim, output_dim) + self.relu = nn.ReLU() + + # Define alignment layers for pruning + self.alignment_layers = [self.fc1, self.fc2, self.fc3] + self.alignment_names = ['layer_0', 'layer_1', 'layer_2'] + self.hidden = {} + + def forward(self, x): + # Store inputs for hooking + batch_size = x.size(0) + x = x.view(batch_size, -1) # Flatten + + # Forward with storing activations + self.hidden['layer_0'] = x + x = self.relu(self.fc1(x)) + + self.hidden['layer_1'] = x + x = self.relu(self.fc2(x)) + + self.hidden['layer_2'] = x + x = self.fc3(x) + + return x + + # Create model based on dataset + model = MLP(input_dim=input_dim, hidden_dim=100, output_dim=10) + model.to(device) + + # Check for alignment layers + if not hasattr(model, 'alignment_layers'): + logger.error("Model doesn't have alignment_layers attribute") + return + + logger.info(f"Created model with {len(model.alignment_layers)} alignment layers") + + # Make a copy for pruning + pruned_model = copy.deepcopy(model) + pruned_model.to(device) + + # Make sure parameters are on device + for param in pruned_model.parameters(): + param.data = param.data.to(device) + + # Evaluate original model + model.eval() + orig_accuracy, orig_loss = dataset.evaluate(model, device) + logger.info(f"Original model: accuracy={orig_accuracy:.2f}%, loss={orig_loss:.4f}") + + # Check zero weights before pruning + logger.info("Weights before pruning:") + orig_zeros, orig_total = count_zero_weights(pruned_model) + + # Setup metric + metric = get_metric('rq') + + # Apply pruning + logger.info(f"Applying pruning: strategy={strategy}, mode={pruning_mode}, " + f"dropout_mode={dropout_mode}, percent={pruning_percent*100:.1f}%") + + # Call progressive_dropout directly + network_accuracies, network_losses = progressive_dropout( + [pruned_model], + dataset, + [0.0, pruning_percent], # Use just two fractions: 0% and the target % + metric, + device, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + strategy=strategy, + show_progress=False + ) + + # Check if pruning was applied + logger.info("Weights after pruning:") + pruned_zeros, pruned_total = count_zero_weights(pruned_model) + + # Calculate how many weights were pruned + weights_pruned = pruned_zeros - orig_zeros + pruned_percent = 100.0 * weights_pruned / pruned_total if pruned_total > 0 else 0 + logger.info(f"Weights pruned: {weights_pruned}/{pruned_total} ({pruned_percent:.2f}%)") + + # Evaluate pruned model + pruned_model.eval() + pruned_accuracy, pruned_loss = dataset.evaluate(pruned_model, device) + logger.info(f"Pruned model: accuracy={pruned_accuracy:.2f}%, loss={pruned_loss:.4f}") + + # Calculate accuracy change + acc_change = pruned_accuracy - orig_accuracy + logger.info(f"Accuracy change: {acc_change:.2f}% points") + + # Summary + logger.info("\nSUMMARY:") + logger.info(f"Strategy: {strategy}") + logger.info(f"Pruning mode: {pruning_mode}") + logger.info(f"Dropout mode: {dropout_mode}") + logger.info(f"Pruning percentage: {pruning_percent*100:.1f}%") + logger.info(f"Original accuracy: {orig_accuracy:.2f}%") + logger.info(f"Pruned accuracy: {pruned_accuracy:.2f}%") + logger.info(f"Accuracy change: {acc_change:.2f}% points") + logger.info(f"Weights pruned: {weights_pruned}/{pruned_total} ({pruned_percent:.2f}%)") + + return { + 'orig_accuracy': orig_accuracy, + 'pruned_accuracy': pruned_accuracy, + 'acc_change': acc_change, + 'weights_pruned': weights_pruned, + 'total_weights': pruned_total, + 'pruned_percent': pruned_percent + } + +def main(): + """Main function.""" + parser = argparse.ArgumentParser(description='Test pruning directly') + parser.add_argument('--model', type=str, default='mlp', help='Model name') + parser.add_argument('--dataset', type=str, default='mnist', help='Dataset name') + parser.add_argument('--strategy', type=str, default='low_rq', + choices=['high_rq', 'low_rq', 'random'], help='Pruning strategy') + parser.add_argument('--pruning-mode', type=str, default='layer_wise', + choices=['global_joint', 'layer_wise', 'layer_isolated'], + help='Pruning mode') + parser.add_argument('--dropout-mode', type=str, default='unscaled', + choices=['scaled', 'unscaled'], help='Dropout mode') + parser.add_argument('--pruning-percent', type=float, default=0.5, + help='Pruning percentage (0.0-1.0)') + parser.add_argument('--device', type=str, default=None, help='Device (cuda/cpu)') + + args = parser.parse_args() + + # Convert percent to fraction if needed + if args.pruning_percent > 1.0: + args.pruning_percent /= 100.0 + + # Run test + verify_pruning( + model_name=args.model, + dataset_name=args.dataset, + strategy=args.strategy, + pruning_mode=args.pruning_mode, + dropout_mode=args.dropout_mode, + pruning_percent=args.pruning_percent, + device=args.device + ) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/_arxiv/archive/scripts/run_benchmark.sh b/_arxiv/archive/scripts/run_benchmark.sh new file mode 100755 index 00000000..db589054 --- /dev/null +++ b/_arxiv/archive/scripts/run_benchmark.sh @@ -0,0 +1,42 @@ +#!/bin/bash + +# Benchmark script for testing tensorized network training +# This script runs multiple benchmark configurations to compare +# training speeds for different numbers of networks + +ROOT_DIR="/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment" +cd $ROOT_DIR + +# Set up common parameters +EPOCHS=1 +DEVICE="cuda" +BATCH_SIZE=64 +HIDDEN_SIZES="32,16" + +echo "----------------------------------------------" +echo "Network Training Method Benchmark" +echo "----------------------------------------------" +echo "Date: $(date)" +echo "Device: $DEVICE" +echo "Hidden sizes: $HIDDEN_SIZES" +echo "Epochs: $EPOCHS" +echo "Batch size: $BATCH_SIZE" +echo "----------------------------------------------" +echo "" + +# Run benchmarks with different network counts +for NUM_NETWORKS in 1 3 5 10 20 +do + echo "=== Running benchmark with $NUM_NETWORKS networks ===" + python benchmark_network_training.py \ + --num_networks $NUM_NETWORKS \ + --hidden_sizes $HIDDEN_SIZES \ + --epochs $EPOCHS \ + --device $DEVICE \ + --batch_size $BATCH_SIZE + echo "" +done + +echo "----------------------------------------------" +echo "Benchmark completed" +echo "----------------------------------------------" \ No newline at end of file diff --git a/_arxiv/archive/scripts/run_cascading_test.py b/_arxiv/archive/scripts/run_cascading_test.py new file mode 100644 index 00000000..5e49e42f --- /dev/null +++ b/_arxiv/archive/scripts/run_cascading_test.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python +""" +Test script for running alignment experiments with cascading layer pruning +""" + +import os +import sys +import logging +import time + +# Add the src directory to the Python path +script_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.join(script_dir, "src")) + +try: + from alignment.config import ExperimentConfig + from alignment.experiments.alignment_experiments import AlignmentExperiment + from alignment.experiments.alignment_experiments import set_logging_level + print("Successfully imported alignment modules") +except ImportError as e: + print(f"Error importing alignment modules: {e}") + sys.exit(1) + +def main(): + # Define experiment parameters + config_file = "configs/config_alignment_experiment.yaml" + experiment_type = "progressive_dropout" + pruning_mode = "cascading_layer" + run_name = "cascading_test" + + print(f"Starting alignment experiment with cascading layer pruning at {time.ctime()}") + start_time = time.time() + + # Set up logging + set_logging_level(logging.INFO) + + # Load configuration + config_path = os.path.join(script_dir, config_file) + config = ExperimentConfig.load(config_path) + + # Set experiment parameters + config.experiment_type = experiment_type + + # Make sure the extra attribute exists + if not hasattr(config, 'extra'): + config.extra = type('ExtraConfig', (), {})() + + config.extra.dropout_pruning_mode = pruning_mode + + # Set experiment name + config.experiment_name = run_name + + # Modify configuration for faster testing + config.training.replicates = 3 # Use fewer networks + config.alignment.dropout_steps = 3 # Use fewer dropout steps + + # Print configuration for debugging + print(f"Using pruning mode: {pruning_mode}") + print(f"Number of replicates: {config.training.replicates}") + print(f"Number of dropout steps: {config.alignment.dropout_steps}") + + # Create and run experiment + experiment = AlignmentExperiment(config) + results, networks = experiment.run() + + end_time = time.time() + print(f"Alignment experiment finished at {time.ctime()}") + print(f"Total duration: {end_time - start_time:.2f} seconds") + + return 0 + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/_arxiv/archive/scripts/run_cascading_test.sh b/_arxiv/archive/scripts/run_cascading_test.sh new file mode 100755 index 00000000..a233ad25 --- /dev/null +++ b/_arxiv/archive/scripts/run_cascading_test.sh @@ -0,0 +1,50 @@ +#!/bin/bash + +# Load necessary modules (if needed) +# module purge +# module load python/3.12.5-fasrc01 +# module load cuda/12.4.1-fasrc01 +# module load cudnn/8.9.2.26_cuda12-fasrc01 + +# Define the experiment parameters +CONFIG_FILE="configs/config_alignment_experiment.yaml" +EXPERIMENT_TYPE="progressive_dropout" +PRUNING_MODE="cascading_layer" +RUN_NAME="cascading_test" + +echo "Starting alignment experiment with cascading layer pruning at $(date)" +start_time=$(date +%s) + +# Add the alignment package to the Python path +export PYTHONPATH=/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment + +# Execute the alignment experiment +cd /n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment +python -c " +import logging +import sys +from alignment.config import ExperimentConfig +from alignment.experiments.alignment_experiments import AlignmentExperiment, set_logging_level + +# Set up logging +set_logging_level(logging.INFO) + +# Load configuration +config_path = '${CONFIG_FILE}' +config = ExperimentConfig.load(config_path) + +# Set experiment parameters +config.experiment_type = '${EXPERIMENT_TYPE}' +config.extra.dropout_pruning_mode = '${PRUNING_MODE}' +config.run.name = '${RUN_NAME}' + +# Create and run experiment +experiment = AlignmentExperiment(config) +results, networks = experiment.run() + +print('Experiment completed successfully') +" + +end_time=$(date +%s) +echo "Alignment experiment finished at $(date)" +echo "Total duration: $((end_time - start_time)) seconds." \ No newline at end of file diff --git a/_arxiv/archive/scripts/run_cascading_with_plots.py b/_arxiv/archive/scripts/run_cascading_with_plots.py new file mode 100755 index 00000000..e4cba3b2 --- /dev/null +++ b/_arxiv/archive/scripts/run_cascading_with_plots.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python +""" +Run alignment experiment with cascading layer pruning and enhanced visualization +""" + +import os +import sys +import logging +import time +import matplotlib.pyplot as plt +import numpy as np +import torch + +# Add the src directory to the Python path +script_dir = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, os.path.join(script_dir, "src")) + +# Import alignment modules +from alignment.config import ExperimentConfig +from alignment.experiments.alignment_experiments import AlignmentExperiment, set_logging_level +from alignment.utils.plotting import plot_dropout_results + +# Set style to match the alignment_preref version +plt.style.use('seaborn-v0_8-whitegrid') +plt.rcParams['figure.figsize'] = (10, 6) +plt.rcParams['lines.linewidth'] = 2.5 +plt.rcParams['axes.grid'] = True +plt.rcParams['grid.alpha'] = 0.3 +plt.rcParams['axes.labelsize'] = 14 +plt.rcParams['xtick.labelsize'] = 12 +plt.rcParams['ytick.labelsize'] = 12 +plt.rcParams['legend.fontsize'] = 12 +plt.rcParams['figure.titlesize'] = 16 + +# Ensure wandb is available +try: + import wandb +except ImportError: + print("Warning: wandb not installed. Install with 'pip install wandb' for full functionality.") + wandb = None + +def enhanced_plotting(results, figure_path, pruning_mode, dropout_mode, experiment_name): + """Enhanced plotting function with consistent colors and styles""" + + # Define consistent colors and styles to match alignment_preref + saved_figures = [] + + # Extract data + dropout_fractions = results["dropout_fractions"] + accuracies = np.array(results["accuracies"]) + + # Plot mean accuracy vs. dropout fraction with enhanced styling + plt.figure(figsize=(12, 8)) + + # Use a richer color palette and add markers for better visibility + plt.plot(dropout_fractions, np.mean(accuracies, axis=0), 'o-', + color='#1f77b4', linewidth=2.5, markersize=8, + label="Mean Accuracy") + + # Add error bands with semi-transparency + plt.fill_between( + dropout_fractions, + np.mean(accuracies, axis=0) - np.std(accuracies, axis=0), + np.mean(accuracies, axis=0) + np.std(accuracies, axis=0), + alpha=0.25, color='#1f77b4' + ) + + # Enhance plot formatting + plt.xlabel("Dropout Fraction", fontsize=14) + plt.ylabel("Accuracy (%)", fontsize=14) + plt.title(f"{experiment_name}: {pruning_mode} Pruning", fontsize=16) + plt.grid(True, alpha=0.3) + plt.ylim(0, 100) # Set consistent y-axis range + plt.xlim(0, max(dropout_fractions) * 1.05) # Add a small margin + plt.legend(loc="upper right", fontsize=12) + + # Add more information to the plot + plt.text(0.02, 0.02, f"Dropout Mode: {dropout_mode}", transform=plt.gca().transAxes, + fontsize=10, bbox=dict(facecolor='white', alpha=0.8)) + + # Save the enhanced figure + if figure_path is not None: + os.makedirs(figure_path, exist_ok=True) + filename = os.path.join( + figure_path, + f"enhanced_{pruning_mode}_{dropout_mode}.png" + ) + plt.savefig(filename, dpi=300, bbox_inches="tight") + saved_figures.append(filename) + + # Log to wandb if available + if wandb is not None and wandb.run is not None: + wandb.log({"dropout_plot": wandb.Image(filename)}) + + plt.close() + else: + plt.show() + + # Also generate the standard plots for comparison + standard_plots = plot_dropout_results( + results, + figure_path, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode, + title_prefix=experiment_name + ) + + saved_figures.extend(standard_plots) + return saved_figures + +def run_experiment(config_path): + """Run the experiment with enhanced visualization""" + print(f"Starting cascading layer pruning experiment at {time.ctime()}") + start_time = time.time() + + # Set up logging + set_logging_level(logging.INFO) + + # Load configuration + config = ExperimentConfig.load(config_path) + + # Ensure we're using progressive dropout with cascading layer + config.experiment_type = "progressive_dropout" + + # Make sure the extra attribute exists + if not hasattr(config, 'extra'): + config.extra = type('ExtraConfig', (), {})() + + # Set the pruning mode to cascading layer + config.extra.dropout_pruning_mode = "cascading_layer" + + # Initialize wandb if available + if wandb is not None and getattr(config.checkpointing, "use_wandb", False): + wandb.init( + project=getattr(config, "wandb_project", "neural_alignment"), + entity=getattr(config, "wandb_entity", None), + name=getattr(config, "experiment_name", "cascading_layer_test"), + config=config.to_dict() + ) + wandb.run.log_code(".", include_fn=lambda path: path.endswith(".py")) + + # Create and run experiment + experiment = AlignmentExperiment(config) + results, networks = experiment.run() + + # Enhance the default plots + if "progressive_dropout" in results: + dropout_results = results["progressive_dropout"] + pruning_mode = config.extra.dropout_pruning_mode + dropout_mode = getattr(config.extra, "dropout_mode", "scaled") + + # Create enhanced plots + saved_figures = enhanced_plotting( + dropout_results, + experiment.figure_path, + pruning_mode, + dropout_mode, + getattr(config, "experiment_name", "Progressive Dropout") + ) + + print(f"Generated {len(saved_figures)} plot files") + + end_time = time.time() + print(f"Experiment finished at {time.ctime()}") + print(f"Total duration: {end_time - start_time:.2f} seconds") + + # Finish wandb run + if wandb is not None and wandb.run is not None: + wandb.finish() + + return results, networks + +if __name__ == "__main__": + if len(sys.argv) > 1: + config_path = sys.argv[1] + else: + config_path = "configs/config_alignment_experiment.yaml" + + results, _ = run_experiment(config_path) \ No newline at end of file diff --git a/_arxiv/archive/scripts/run_fixed_experiment.py b/_arxiv/archive/scripts/run_fixed_experiment.py new file mode 100644 index 00000000..2935332c --- /dev/null +++ b/_arxiv/archive/scripts/run_fixed_experiment.py @@ -0,0 +1,402 @@ +#!/usr/bin/env python +""" +Run a fixed version of the alignment experiment. +""" + +import os +import sys +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +import numpy as np +import logging +from tqdm import tqdm + +# Configure basic logging +logging.basicConfig(level=logging.INFO, + format='%(asctime)s [%(levelname)s] %(message)s', + handlers=[logging.StreamHandler()]) +logger = logging.getLogger(__name__) + +# Add the project to path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "."))) + +# Import from the alignment package +from alignment.config import ExperimentConfig +from alignment.models.registry import create_model +from alignment.datasets import load_dataset +from alignment.metrics import get_metric +from alignment.utils.plotting import plot_dropout_results + +def train_networks(networks, dataset, device, num_epochs=5, learning_rate=0.001): + """Train multiple networks on the given dataset.""" + logger.info(f"Training {len(networks)} networks for {num_epochs} epochs") + + # Track training history for plotting + training_history = { + 'train_loss': [], + 'train_acc': [], + 'test_loss': [], + 'test_acc': [] + } + + # Train each network + trained_networks = [] + for i, network in enumerate(networks): + logger.info(f"Training network {i+1}/{len(networks)}") + + # Move network to the device + network = network.to(device) + + # Create optimizer + optimizer = optim.Adam(network.parameters(), lr=learning_rate) + + # Training loop + network.train() + history = { + 'train_loss': [], + 'train_acc': [], + 'test_loss': [], + 'test_acc': [] + } + + for epoch in range(num_epochs): + running_loss = 0.0 + correct = 0 + total = 0 + + # Train on each batch + for inputs, targets in tqdm(dataset.train_loader, desc=f"Network {i+1}, Epoch {epoch+1}/{num_epochs}"): + inputs, targets = inputs.to(device), targets.to(device) + + # Zero the parameter gradients + optimizer.zero_grad() + + # Forward pass + outputs = network(inputs) + loss = F.cross_entropy(outputs, targets) + + # Backward pass and optimize + loss.backward() + optimizer.step() + + # Track statistics + running_loss += loss.item() + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + # Calculate epoch statistics + train_loss = running_loss / len(dataset.train_loader) + train_acc = 100.0 * correct / total + + # Evaluate on test set + network.eval() + test_correct = 0 + test_total = 0 + test_loss_sum = 0.0 + + with torch.no_grad(): + for inputs, targets in dataset.test_loader: + inputs, targets = inputs.to(device), targets.to(device) + outputs = network(inputs) + + # Calculate loss + loss = F.cross_entropy(outputs, targets, reduction='sum') + test_loss_sum += loss.item() + + # Calculate accuracy + _, predicted = outputs.max(1) + test_total += targets.size(0) + test_correct += predicted.eq(targets).sum().item() + + test_acc = 100.0 * test_correct / test_total + test_loss = test_loss_sum / test_total + + # Store history + history['train_loss'].append(train_loss) + history['train_acc'].append(train_acc) + history['test_loss'].append(test_loss) + history['test_acc'].append(test_acc) + + # Log progress + logger.info(f"Network {i+1}, Epoch {epoch+1}/{num_epochs}: " + f"Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%, " + f"Test Loss={test_loss:.4f}, Test Acc={test_acc:.2f}%") + + # Add to trained networks + trained_networks.append(network) + + # Accumulate training history (average across networks) + if i == 0: + # First network, initialize history + training_history = history + else: + # Average with previous networks + for key in training_history: + if key in history: + # Calculate running average + for epoch_idx in range(len(history[key])): + if epoch_idx < len(training_history[key]): + training_history[key][epoch_idx] = (training_history[key][epoch_idx] * i + history[key][epoch_idx]) / (i + 1) + + logger.info(f"Completed training {len(networks)} networks") + return trained_networks, training_history + +def test_pruning_strategies(networks, dataset, dropout_fractions, device): + """Test different pruning strategies on trained networks.""" + logger.info(f"Testing pruning strategies on {len(networks)} networks") + + # Initialize results + results = { + 'dropout_fractions': dropout_fractions, + 'accuracies': {'high_rq': [], 'low_rq': [], 'random': []}, + 'stds': {'high_rq': [], 'low_rq': [], 'random': []}, + 'losses': {'high_rq': [], 'low_rq': [], 'random': []} + } + + # For each pruning percentage, test all networks + for prune_idx, prune_percent in enumerate(dropout_fractions): + logger.info(f"Testing pruning percentage: {prune_percent*100:.1f}%") + + # Store results for each network at this pruning percentage + strategy_results = { + 'high_rq': {'acc': [], 'loss': []}, + 'low_rq': {'acc': [], 'loss': []}, + 'random': {'acc': [], 'loss': []} + } + + # Process each network + for net_idx, network in enumerate(networks): + logger.info(f"Processing network {net_idx+1}/{len(networks)}") + + # Save original weights + original_weights = {} + original_biases = {} + + for i, layer in enumerate(network.alignment_layers): + if hasattr(layer, "weight") and layer.weight is not None: + original_weights[i] = layer.weight.data.clone() + if hasattr(layer, "bias") and layer.bias is not None: + original_biases[i] = layer.bias.data.clone() + + # Get original accuracy + if prune_idx == 0 or net_idx == 0: + network.eval() + correct = 0 + total = 0 + total_loss = 0.0 + + with torch.no_grad(): + for inputs, targets in dataset.test_loader: + inputs, targets = inputs.to(device), targets.to(device) + outputs = network(inputs) + + # Compute loss + loss = F.cross_entropy(outputs, targets, reduction='sum') + total_loss += loss.item() + + # Compute accuracy + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + orig_accuracy = 100.0 * correct / total + orig_loss = total_loss / total + + logger.info(f"Original accuracy: {orig_accuracy:.2f}%, loss: {orig_loss:.4f}") + + # Test each strategy + for strategy in ["high_rq", "low_rq", "random"]: + # Restore original weights + for i, layer in enumerate(network.alignment_layers): + if i in original_weights: + layer.weight.data = original_weights[i].clone() + if i in original_biases and hasattr(layer, "bias") and layer.bias is not None: + layer.bias.data = original_biases[i].clone() + + # Skip pruning for 0% case + if prune_percent == 0.0: + acc = orig_accuracy + loss = orig_loss + else: + # Apply pruning based on strategy + total_neurons = 0 + total_pruned = 0 + + # Apply pruning to each layer + for i, layer in enumerate(network.alignment_layers): + if i not in original_weights: + continue + + # Compute neuron importance scores (using weight magnitude as proxy) + weights = layer.weight.data + input_dim = weights.shape[1] + total_neurons += input_dim + + # Calculate importance scores (weight magnitude) + neuron_scores = [torch.norm(weights[:, j]).item() for j in range(input_dim)] + + # Calculate how many neurons to prune + num_to_drop = max(1, int(input_dim * prune_percent)) if prune_percent > 0 else 0 + total_pruned += num_to_drop + + if num_to_drop > 0: + # Get indices to drop based on strategy + if strategy == "high_rq": # Drop highest alignment neurons + sorted_indices = np.argsort(neuron_scores)[::-1] # Sort descending + to_drop = sorted_indices[:num_to_drop] + elif strategy == "low_rq": # Drop lowest alignment neurons + sorted_indices = np.argsort(neuron_scores) # Sort ascending + to_drop = sorted_indices[:num_to_drop] + else: # Random pruning + all_indices = list(range(input_dim)) + np.random.shuffle(all_indices) + to_drop = all_indices[:num_to_drop] + + # Zero out weights for these neurons + for idx in to_drop: + if idx < weights.shape[1]: + layer.weight.data[:, idx] = 0.0 + if hasattr(layer, "bias") and layer.bias is not None and idx < layer.bias.data.shape[0]: + layer.bias.data[idx] = 0.0 + + logger.info(f"Pruned {total_pruned}/{total_neurons} neurons with strategy: {strategy}") + + # Evaluate pruned network + network.eval() + correct = 0 + total = 0 + total_loss = 0.0 + + with torch.no_grad(): + for inputs, targets in dataset.test_loader: + inputs, targets = inputs.to(device), targets.to(device) + outputs = network(inputs) + + # Compute loss + loss = F.cross_entropy(outputs, targets, reduction='sum') + total_loss += loss.item() + + # Compute accuracy + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + acc = 100.0 * correct / total + loss = total_loss / total + + # Store results for this network and strategy + strategy_results[strategy]['acc'].append(acc) + strategy_results[strategy]['loss'].append(loss) + + logger.info(f"Strategy {strategy}, Accuracy: {acc:.2f}%, Loss: {loss:.4f}") + + # Calculate mean and std for each strategy at this pruning percentage + for strategy in ["high_rq", "low_rq", "random"]: + # Get all networks' results for this strategy and pruning percentage + accs = strategy_results[strategy]['acc'] + losses = strategy_results[strategy]['loss'] + + # Calculate statistics + mean_acc = np.mean(accs) + std_acc = np.std(accs) + mean_loss = np.mean(losses) + + # Add to results + results['accuracies'][strategy].append(mean_acc) + results['stds'][strategy].append(std_acc) + results['losses'][strategy].append(mean_loss) + + logger.info(f"Pruning {prune_percent*100:.1f}%, Strategy {strategy}: " + f"Mean acc: {mean_acc:.2f}%, Std: {std_acc:.2f}%") + + return results + +def main(): + """Main function to run the experiment.""" + # Set random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Load configuration + config_path = "configs/config_alignment_experiment.yaml" + config = ExperimentConfig.load(config_path) + + # Set device + device = "cuda" if torch.cuda.is_available() else "cpu" + logger.info(f"Using device: {device}") + + # Load dataset + logger.info(f"Loading dataset: {config.dataset.dataset_name}") + dataset = load_dataset(config.dataset) + + # Create multiple networks (replicates) + num_networks = min(getattr(config.training, "replicates", 5), 2) # Limit to 2 networks for faster testing + logger.info(f"Creating {num_networks} networks") + + networks = [] + for i in range(num_networks): + logger.info(f"Creating model {i+1}/{num_networks}: {config.model.model_name}") + network = create_model(config.model) + networks.append(network) + + # Get dropout fractions from config + dropin_min = config.alignment.dropout_min + dropin_max = config.alignment.dropout_max + num_dropout_fractions = config.alignment.dropout_steps + dropout_fractions = np.linspace(dropin_min, dropin_max, num_dropout_fractions).tolist() + + # Train all networks + trained_networks, training_history = train_networks( + networks, + dataset, + device, + num_epochs=getattr(config.training, "epochs", 5), + learning_rate=getattr(config.training, "learning_rate", 0.001) + ) + + # Test pruning strategies + pruning_results = test_pruning_strategies( + trained_networks, + dataset, + dropout_fractions, + device + ) + + # Add training history to results + pruning_results['training_history'] = training_history + + # Create output directory + results_dir = "debug_output" + os.makedirs(results_dir, exist_ok=True) + + # Generate plots + try: + logger.info("Generating pruning plots with error bars") + + # Get pruning mode and dropout mode from config + pruning_mode = getattr(config.extra, "dropout_pruning_mode", "global_joint") + dropout_mode = getattr(config.extra, "dropout_mode", "scaled") + + saved_plots = plot_dropout_results( + pruning_results, + results_dir, + title_prefix="Fixed Pruning Test", + pruning_mode=pruning_mode, + dropout_mode=dropout_mode + ) + + if saved_plots: + logger.info(f"Generated {len(saved_plots)} plots: {saved_plots}") + + except Exception as e: + logger.error(f"Error generating dropout plots: {str(e)}") + import traceback + logger.error(traceback.format_exc()) + + logger.info("Experiment completed successfully!") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/_arxiv/archive/scripts/run_multi_strategy_experiment.py b/_arxiv/archive/scripts/run_multi_strategy_experiment.py new file mode 100644 index 00000000..95d2a8cc --- /dev/null +++ b/_arxiv/archive/scripts/run_multi_strategy_experiment.py @@ -0,0 +1,614 @@ +#!/usr/bin/env python +""" +Run an experiment with multiple pruning strategies. + +This script executes neural network pruning using three different strategies: +1. high_rq: Prune neurons with highest alignment scores (weight magnitudes) +2. low_rq: Prune neurons with lowest alignment scores (weight magnitudes) +3. random: Prune neurons randomly + +The results are plotted for comparison. +""" + +import os +import sys +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +import numpy as np +import logging +from tqdm import tqdm +import matplotlib.pyplot as plt +import copy + +# Configure basic logging +logging.basicConfig(level=logging.INFO, + format='%(asctime)s [%(levelname)s] %(message)s', + handlers=[logging.StreamHandler()]) +logger = logging.getLogger(__name__) + +# Add the project to path +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "."))) + +# Import from the alignment package +from alignment.config import ExperimentConfig +from alignment.models.registry import create_model +from alignment.datasets import load_dataset +from alignment.metrics import get_metric +from alignment.utils.plotting import plot_dropout_results + +def train_networks(networks, dataset, device, num_epochs=5, learning_rate=0.001): + """Train multiple networks efficiently in one training loop.""" + logger.info(f"Training {len(networks)} networks for {num_epochs} epochs") + + # Track training history for plotting + training_history = { + 'train_loss': [], + 'train_acc': [], + 'test_loss': [], + 'test_acc': [] + } + + # Set up optimizer for each network + optimizers = [] + for network in networks: + network.to(device) + optimizers.append(torch.optim.Adam(network.parameters(), lr=learning_rate)) + + # Train all networks for each epoch + epoch_pbar = tqdm(range(num_epochs), desc="Training epochs", position=0) + for epoch in epoch_pbar: + # Initialize epoch stats + epoch_train_loss = 0.0 + epoch_train_acc = 0.0 + epoch_test_loss = 0.0 + epoch_test_acc = 0.0 + + # Training phase + net_pbar = tqdm(enumerate(zip(networks, optimizers)), + desc=f"Epoch {epoch+1}/{num_epochs} networks", + total=len(networks), + position=1, + leave=False) + + for network_idx, (network, optimizer) in net_pbar: + network.train() + running_loss = 0.0 + correct = 0 + total = 0 + + # Train on each batch + batch_pbar = tqdm(dataset.train_loader, + desc=f"Network {network_idx+1}/{len(networks)}", + position=2, + leave=False) + + for inputs, targets in batch_pbar: + inputs, targets = inputs.to(device), targets.to(device) + + # Zero the parameter gradients + optimizer.zero_grad() + + # Forward pass + outputs = network(inputs) + loss = F.cross_entropy(outputs, targets) + + # Backward pass and optimize + loss.backward() + optimizer.step() + + # Track statistics + running_loss += loss.item() + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + # Update batch progress bar + if total > 0: + batch_pbar.set_postfix({ + 'loss': running_loss / len(batch_pbar), + 'acc': 100.0 * correct / total + }) + + # Calculate network training statistics + if total > 0: + network_train_loss = running_loss / len(dataset.train_loader) + network_train_acc = 100.0 * correct / total + + # Accumulate for epoch average + epoch_train_loss += network_train_loss + epoch_train_acc += network_train_acc + + # Update network progress bar + net_pbar.set_postfix({ + 'train_loss': network_train_loss, + 'train_acc': network_train_acc + }) + + # Evaluation phase + network.eval() + test_correct = 0 + test_total = 0 + test_loss_sum = 0.0 + + # Evaluation progress bar + eval_pbar = tqdm(dataset.test_loader, + desc=f"Evaluating network {network_idx+1}", + position=2, + leave=False) + + with torch.no_grad(): + for inputs, targets in eval_pbar: + inputs, targets = inputs.to(device), targets.to(device) + outputs = network(inputs) + + # Calculate loss + loss = F.cross_entropy(outputs, targets, reduction='sum') + test_loss_sum += loss.item() + + # Calculate accuracy + _, predicted = outputs.max(1) + test_total += targets.size(0) + test_correct += predicted.eq(targets).sum().item() + + # Update eval progress bar + if test_total > 0: + eval_pbar.set_postfix({ + 'test_loss': test_loss_sum / test_total, + 'test_acc': 100.0 * test_correct / test_total + }) + + # Calculate network testing statistics + if test_total > 0: + network_test_loss = test_loss_sum / test_total + network_test_acc = 100.0 * test_correct / test_total + + # Accumulate for epoch average + epoch_test_loss += network_test_loss + epoch_test_acc += network_test_acc + + # Calculate and store epoch averages + epoch_train_loss /= len(networks) + epoch_train_acc /= len(networks) + epoch_test_loss /= len(networks) + epoch_test_acc /= len(networks) + + training_history['train_loss'].append(epoch_train_loss) + training_history['train_acc'].append(epoch_train_acc) + training_history['test_loss'].append(epoch_test_loss) + training_history['test_acc'].append(epoch_test_acc) + + # Update epoch progress bar + epoch_pbar.set_postfix({ + 'train_loss': epoch_train_loss, + 'train_acc': f"{epoch_train_acc:.2f}%", + 'test_loss': epoch_test_loss, + 'test_acc': f"{epoch_test_acc:.2f}%" + }) + + # Log progress + logger.info(f"Epoch {epoch+1}/{num_epochs}: " + f"Train Loss={epoch_train_loss:.4f}, Train Acc={epoch_train_acc:.2f}%, " + f"Test Loss={epoch_test_loss:.4f}, Test Acc={epoch_test_acc:.2f}%") + + logger.info(f"Completed training {len(networks)} networks.") + return networks, training_history + +def run_pruning(networks, dataset, dropout_fractions, device, pruning_mode="layer_wise", dropout_mode="scaled"): + """Run pruning with three different strategies and collect results in a more efficient way.""" + # Define the strategies to evaluate + strategies = ["high_rq", "low_rq", "random"] + + # Initialize results structure + results = { + 'dropout_fractions': dropout_fractions, + 'accuracies': {strategy: [] for strategy in strategies}, + 'stds': {strategy: [] for strategy in strategies}, + 'losses': {strategy: [] for strategy in strategies} + } + + # Store original weights and biases for all networks + logger.info("Saving original network weights") + original_weights = {} + original_biases = {} + + for net_idx, network in enumerate(networks): + original_weights[net_idx] = {} + original_biases[net_idx] = {} + + for layer_idx, layer in enumerate(network.alignment_layers): + if hasattr(layer, "weight") and layer.weight is not None: + original_weights[net_idx][layer_idx] = layer.weight.data.clone() + if hasattr(layer, "bias") and layer.bias is not None: + original_biases[net_idx][layer_idx] = layer.bias.data.clone() + + # First, get original accuracy for all networks (0% pruning) + logger.info("Evaluating original network performance (0% pruning)") + orig_accs = [] + orig_losses = [] + + # Move all networks to device + for network in networks: + network.to(device) + + # Evaluate all networks at baseline + for network in networks: + network.eval() + correct = 0 + total = 0 + total_loss = 0.0 + + with torch.no_grad(): + for inputs, targets in dataset.test_loader: + inputs, targets = inputs.to(device), targets.to(device) + outputs = network(inputs) + + # Compute loss + loss = F.cross_entropy(outputs, targets, reduction='sum') + total_loss += loss.item() + + # Compute accuracy + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + if total > 0: + orig_accs.append(100.0 * correct / total) + orig_losses.append(total_loss / total) + + # Calculate baseline statistics + orig_mean_acc = np.mean(orig_accs) + orig_std_acc = np.std(orig_accs) + orig_mean_loss = np.mean(orig_losses) + + # Add baseline values to all strategies (they start from the same point) + for strategy in strategies: + results['accuracies'][strategy].append(orig_mean_acc) + results['stds'][strategy].append(orig_std_acc) + results['losses'][strategy].append(orig_mean_loss) + + logger.info(f"Original accuracy: {orig_mean_acc:.2f}% ± {orig_std_acc:.2f}%") + + # Iterate through pruning fractions (skip the first one which is 0%) + fraction_pbar = tqdm(enumerate(dropout_fractions), + total=len(dropout_fractions), + desc="Pruning fractions", + position=0) + + for fraction_idx, fraction in fraction_pbar: + if fraction_idx == 0 and fraction == 0.0: + continue + + fraction_pbar.set_description(f"Pruning fraction: {fraction:.2f}") + + # For each strategy, process all networks at once + strategy_pbar = tqdm(strategies, desc=f"Pruning strategies", position=1, leave=False) + + for strategy in strategy_pbar: + strategy_pbar.set_description(f"Strategy: {strategy}") + + # Restore original weights for all networks at once + for net_idx, network in enumerate(networks): + for layer_idx in original_weights[net_idx]: + layer = network.alignment_layers[layer_idx] + layer.weight.data = original_weights[net_idx][layer_idx].clone() + if layer_idx in original_biases[net_idx] and hasattr(layer, "bias") and layer.bias is not None: + layer.bias.data = original_biases[net_idx][layer_idx].clone() + + # Apply pruning to all networks in parallel + for net_idx, network in enumerate(networks): + # Apply pruning based on the mode + if pruning_mode == "global_joint": + # Collect all neurons and their scores across all layers + all_neurons = [] + + for layer_idx, layer in enumerate(network.alignment_layers): + if layer_idx not in original_weights[net_idx]: + continue + + weights = layer.weight.data + input_dim = weights.shape[1] + + # Calculate importance scores + neuron_scores = torch.norm(weights, dim=0).cpu().numpy() + + # Store (layer_idx, neuron_idx, score) tuples + for j, score in enumerate(neuron_scores): + all_neurons.append((layer_idx, j, score)) + + # Sort neurons by score based on strategy + if strategy == "high_rq": + # Sort by highest scores first + all_neurons.sort(key=lambda x: x[2], reverse=True) + elif strategy == "low_rq": + # Sort by lowest scores first + all_neurons.sort(key=lambda x: x[2]) + elif strategy == "random": + # Shuffle neurons randomly + import random + random.shuffle(all_neurons) + + # Calculate how many neurons to prune + total_neurons = len(all_neurons) + num_to_drop = int(total_neurons * fraction) + + if num_to_drop > 0: + # Get indices to drop + to_drop = all_neurons[:num_to_drop] + + # Apply pruning + for layer_idx, neuron_idx, _ in to_drop: + layer = network.alignment_layers[layer_idx] + # Zero out weights for this neuron + if neuron_idx < layer.weight.data.shape[1]: + layer.weight.data[:, neuron_idx] = 0.0 + if hasattr(layer, "bias") and layer.bias is not None and neuron_idx < layer.bias.data.shape[0]: + layer.bias.data[neuron_idx] = 0.0 + + elif pruning_mode == "layer_wise": + # Apply pruning to each layer individually - can optimize with tensor operations + for layer_idx, layer in enumerate(network.alignment_layers): + if layer_idx not in original_weights[net_idx]: + continue + + weights = layer.weight.data + input_dim = weights.shape[1] + + # Calculate importance scores - vectorized version + neuron_scores = torch.norm(weights, dim=0).cpu().numpy() + + # Calculate how many neurons to prune in this layer + num_to_drop = int(input_dim * fraction) + + if num_to_drop > 0: + # Get neurons to drop based on strategy + if strategy == "high_rq": + # Sort by highest scores first (descending) + sorted_indices = np.argsort(neuron_scores)[::-1] + to_drop = sorted_indices[:num_to_drop] + elif strategy == "low_rq": + # Sort by lowest scores first (ascending) + sorted_indices = np.argsort(neuron_scores) + to_drop = sorted_indices[:num_to_drop] + elif strategy == "random": + # Choose random neurons + all_indices = np.arange(input_dim) + np.random.shuffle(all_indices) + to_drop = all_indices[:num_to_drop] + + # Apply pruning + if to_drop.size > 0: + # Create a mask tensor + mask = torch.ones_like(weights) + + # Convert numpy array to tensor and ensure it's properly formatted for indexing + to_drop_tensor = torch.tensor(to_drop.tolist(), device=weights.device, dtype=torch.long) + + # Use tensor indexing + mask[:, to_drop_tensor] = 0.0 + + # Apply mask to weights + layer.weight.data = weights * mask + + # Apply to bias if it exists + if hasattr(layer, "bias") and layer.bias is not None: + bias_mask = torch.ones_like(layer.bias.data) + valid_indices = [idx for idx in to_drop if idx < layer.bias.data.shape[0]] + if valid_indices: + bias_indices = torch.tensor(valid_indices, device=layer.bias.device, dtype=torch.long) + bias_mask[bias_indices] = 0.0 + layer.bias.data = layer.bias.data * bias_mask + + # Evaluate all pruned networks efficiently + network_accs = [] + network_losses = [] + + for network in networks: + # Evaluate + network.eval() + correct = 0 + total = 0 + total_loss = 0.0 + + with torch.no_grad(): + for inputs, targets in dataset.test_loader: + inputs, targets = inputs.to(device), targets.to(device) + outputs = network(inputs) + + # Compute loss + loss = F.cross_entropy(outputs, targets, reduction='sum') + total_loss += loss.item() + + # Compute accuracy + _, predicted = outputs.max(1) + total += targets.size(0) + correct += predicted.eq(targets).sum().item() + + if total > 0: + accuracy = 100.0 * correct / total + loss = total_loss / total + + network_accs.append(accuracy) + network_losses.append(loss) + + # Calculate statistics for this fraction and strategy + if network_accs: + mean_acc = np.mean(network_accs) + std_acc = np.std(network_accs) + mean_loss = np.mean(network_losses) + + # Add to results + results['accuracies'][strategy].append(mean_acc) + results['stds'][strategy].append(std_acc) + results['losses'][strategy].append(mean_loss) + + logger.info(f" {strategy}: Accuracy {mean_acc:.2f}% ± {std_acc:.2f}%") + + # Restore original weights for all networks + logger.info("Restoring original weights") + for net_idx, network in enumerate(networks): + for layer_idx in original_weights[net_idx]: + layer = network.alignment_layers[layer_idx] + layer.weight.data = original_weights[net_idx][layer_idx].clone() + if layer_idx in original_biases[net_idx] and hasattr(layer, "bias") and layer.bias is not None: + layer.bias.data = original_biases[net_idx][layer_idx].clone() + + return results + +def plot_results(results, pruning_mode, dropout_mode, output_dir="multi_strategy_output"): + """Create plots showing the results of different pruning strategies.""" + # Ensure output directory exists + os.makedirs(output_dir, exist_ok=True) + + # Extract data for plotting + fractions = results['dropout_fractions'] + strategies = list(results['accuracies'].keys()) + + # Set up plot colors and markers + colors = {'high_rq': 'red', 'low_rq': 'green', 'random': 'blue'} + markers = {'high_rq': 'o', 'low_rq': 's', 'random': '^'} + labels = { + 'high_rq': 'Prune Highest Magnitude', + 'low_rq': 'Prune Lowest Magnitude', + 'random': 'Prune Random' + } + + # Create accuracy plot + plt.figure(figsize=(10, 6)) + + for strategy in strategies: + accs = results['accuracies'][strategy] + stds = results['stds'][strategy] + + plt.errorbar( + fractions, + accs, + yerr=stds, + label=labels[strategy], + color=colors[strategy], + marker=markers[strategy], + capsize=4, + markersize=8, + linewidth=2 + ) + + plt.title(f'Accuracy vs. Pruning Fraction\n({pruning_mode} pruning, {dropout_mode} mode)', fontsize=14) + plt.xlabel('Pruning Fraction', fontsize=12) + plt.ylabel('Accuracy (%)', fontsize=12) + plt.grid(True, linestyle='--', alpha=0.7) + plt.legend(fontsize=12) + plt.tight_layout() + + # Save plot + accuracy_plot_file = os.path.join(output_dir, f'accuracy_plot_{pruning_mode}.png') + plt.savefig(accuracy_plot_file, dpi=300) + plt.close() + + # Also use the alignment plotting utility if available + try: + from alignment.utils.plotting import plot_dropout_results + + saved_plots = plot_dropout_results( + results, + output_dir, + title_prefix=f"Multi-Strategy Pruning Test", + pruning_mode=pruning_mode, + dropout_mode=dropout_mode + ) + + if saved_plots: + logger.info(f"Generated {len(saved_plots)} plots using alignment plotting utility") + return saved_plots + except Exception as e: + logger.warning(f"Could not use alignment plotting utility: {str(e)}") + + return [accuracy_plot_file] + +def main(): + """Main function to run the multi-strategy pruning experiment.""" + # Set random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + # Load configuration + config_path = "configs/config_alignment_experiment.yaml" + config = ExperimentConfig.load(config_path) + + # Set device + device = "cuda" if torch.cuda.is_available() else "cpu" + logger.info(f"Using device: {device}") + + # Load dataset + logger.info(f"Loading dataset: {config.dataset.dataset_name}") + dataset = load_dataset(config.dataset) + + # Create multiple networks (replicates) + num_networks = min(getattr(config.training, "replicates", 5), 3) # Limit to 3 networks for faster testing + logger.info(f"Creating {num_networks} networks") + + networks = [] + for i in range(num_networks): + logger.info(f"Creating model {i+1}/{num_networks}: {config.model.model_name}") + network = create_model(config.model) + networks.append(network) + + # Get dropout fractions from config + dropin_min = config.alignment.dropout_min + dropin_max = config.alignment.dropout_max + num_dropout_fractions = config.alignment.dropout_steps + dropout_fractions = np.linspace(dropin_min, dropin_max, num_dropout_fractions).tolist() + + # Get pruning and dropout modes + pruning_mode = getattr(config.extra, "dropout_pruning_mode", "layer_wise") + dropout_mode = getattr(config.extra, "dropout_mode", "scaled") + + # Train all networks + trained_networks, training_history = train_networks( + networks, + dataset, + device, + num_epochs=getattr(config.training, "epochs", 10), + learning_rate=getattr(config.training, "learning_rate", 0.001) + ) + + # Run pruning with different strategies + pruning_results = run_pruning( + trained_networks, + dataset, + dropout_fractions, + device, + pruning_mode=pruning_mode, + dropout_mode=dropout_mode + ) + + # Add training history to results + pruning_results['training_history'] = training_history + + # Create output directory + results_dir = "multi_strategy_output" + os.makedirs(results_dir, exist_ok=True) + + # Generate plots + try: + logger.info("Generating multi-strategy pruning plots") + + saved_plots = plot_results( + pruning_results, + pruning_mode, + dropout_mode, + output_dir=results_dir + ) + + if saved_plots: + logger.info(f"Generated {len(saved_plots)} plots: {saved_plots}") + + except Exception as e: + logger.error(f"Error generating plots: {str(e)}") + import traceback + logger.error(traceback.format_exc()) + + logger.info("Multi-strategy experiment completed successfully!") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/_arxiv/archive/scripts/teacher_student.py b/_arxiv/archive/scripts/teacher_student.py new file mode 100644 index 00000000..3109341d --- /dev/null +++ b/_arxiv/archive/scripts/teacher_student.py @@ -0,0 +1,93 @@ +from tqdm import tqdm +import numpy as np +import torch + +from matplotlib import pyplot as plt + +import os +import sys + +mainPath = os.path.dirname(os.path.abspath(__file__)) + "/.." +sys.path.append(mainPath) + +from networkAlignmentAnalysis.models.registry import get_model +from networkAlignmentAnalysis import utils +from networkAlignmentAnalysis import train + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + +def get_data(modes, signal_dist, noise_amplitude, batch_size): + signal = signal_dist.sample((batch_size,)) + data = signal @ modes + torch.normal(0, noise_amplitude, (batch_size, modes.size(1))) + return signal, data + + +def train(teacher, student, modes, signal_dist, noise_amplitude, batch_size=256, num_epochs=1000): + optimizer = torch.optim.SGD(student.parameters(), lr=1e-2) + lossfn = torch.nn.MSELoss() + track_loss = [] + alignment = [] + for epoch in tqdm(range(num_epochs)): + signal, data = get_data(modes, signal_dist, noise_amplitude, batch_size) + signal, data = signal.to(DEVICE), data.to(DEVICE) + + optimizer.zero_grad() + target = teacher(signal) + output = student(data, store_hidden=True) + loss = lossfn(output, target) + loss.backward() + optimizer.step() + track_loss.append(loss.item()) + alignment.append(student.measure_alignment(data, precomputed=True)) + + alignment = [torch.stack(a) for a in utils.transpose_list(alignment)] # list across layers, (epochs x dim) + return student, track_loss, alignment + + +if __name__ == "__main__": + print("using device: ", DEVICE) + + N = 100 + num_modes = 50 + shape = 0.7 * torch.ones((num_modes,)) + modes = torch.normal(0, 1, (num_modes, N)) + amplitude = torch.rand((num_modes,)) * 0.9 + 0.1 + signal_dist = torch.distributions.gamma.Gamma(shape, amplitude) + noise_amplitude = 0.2 + + model_name = "MLP" + teacher = get_model(model_name, build=True, input_dim=num_modes, dropout=0.0, ignore_flag=False, linear=True).to(DEVICE).eval() + student = get_model(model_name, build=True, input_dim=N, dropout=0.0, ignore_flag=False).to(DEVICE) + + student, track_loss, alignment = train(teacher, student, modes, signal_dist, noise_amplitude) + mean_alignment = torch.stack([torch.mean(align, dim=1) for align in alignment]) + + fig, ax = plt.subplots(1, 3, figsize=(9, 3), layout="constrained") + ax[0].plot(track_loss) + ax[0].set_ylim(0, 0.5) + for layer in range(len(alignment)): + ax[1].plot(mean_alignment[layer], label=f"layer{layer}") + ax[1].legend() + + signal, data = get_data(modes, signal_dist, noise_amplitude, 1) + ax[2].scatter(teacher(signal).detach(), student(data).detach()) + plt.show() + + N, S = 2, 10000 + num_modes = 2 + nonnormality = 1.5 + a = b = 1 / nonnormality + modes = np.random.normal(0, 1, (num_modes, N)) * np.random.normal(0, 1, num_modes).reshape(-1, 1) + signal = np.random.gamma(0.5, 1, (S, num_modes)) * np.sign(np.random.random((S, num_modes)) - 0.5) + data = signal @ modes + w, v = [np.array(o) for o in utils.smart_pca(torch.tensor(data).T)] + print(w.shape, v.shape) + + plt.scatter(data[:, 0], data[:, 1], s=1, c=("b", 0.4)) + for mode in modes: + plt.plot([0, mode[0]], [0, mode[1]], c="r") + for evec in v.T: + plt.plot([0, evec[0]], [0, evec[1]], c="k") + + plt.show() diff --git a/_arxiv/archive/scripts/timetests/test_batched_alignment.py b/_arxiv/archive/scripts/timetests/test_batched_alignment.py new file mode 100644 index 00000000..738ce96c --- /dev/null +++ b/_arxiv/archive/scripts/timetests/test_batched_alignment.py @@ -0,0 +1,41 @@ +import time +import torch +from alignment import utils + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +# B, D, S, C = 1024, 25, 784, 32 +B, D, S, C = 1024, 800, 196, 64 +input = torch.normal(0, 1, (B, D, S)).to(DEVICE) +weight = torch.normal(0, 1, (C, D)).to(DEVICE) + + +def get_align_usual(input, weight): + var_stride = torch.mean(torch.var(input, dim=1), dim=0) + align_stride = torch.stack([utils.alignment(input[:, :, i], weight) for i in range(S)], dim=1) + return utils.weighted_average(align_stride, var_stride.view(1, -1), 1, ignore_nan=True) + + +def get_align_new(input, weight): + var_stride = torch.mean(torch.var(input, dim=1), dim=0) + cc = utils.batch_cov(input.transpose(0, 2)) + rq = torch.sum(torch.matmul(weight, cc) * weight, axis=2) / torch.sum(weight * weight, axis=1) + prq = rq / torch.diagonal(cc, dim1=1, dim2=2).sum(1, keepdim=True) + return utils.weighted_average(prq, var_stride.view(-1, 1), 0, ignore_nan=True) + + +au = get_align_usual(input, weight) +an = get_align_new(input, weight) +print("align_usual and align_new produce same result?", torch.allclose(au, an)) + +num_tests = 100 + +t = time.time() +for _ in range(num_tests): + _ = get_align_usual(input, weight) +print("get_align_usual : time per test=", (time.time() - t) / num_tests) + +t = time.time() +for _ in range(num_tests): + _ = get_align_new(input, weight) +print("get_align_new : time per test=", (time.time() - t) / num_tests) diff --git a/_arxiv/archive/scripts/timetests/test_dataloader.py b/_arxiv/archive/scripts/timetests/test_dataloader.py new file mode 100644 index 00000000..130c2070 --- /dev/null +++ b/_arxiv/archive/scripts/timetests/test_dataloader.py @@ -0,0 +1,106 @@ +import os +import sys + +mainPath = os.path.dirname(os.path.abspath(__file__)) + "/.." +sys.path.append(mainPath) + +import time +from tqdm import tqdm + +import torch +import torchvision +from torchvision.transforms import v2 as transforms +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim + +from alignment import files + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + +class Net(nn.Module): + def __init__(self): + super().__init__() + self.conv1 = nn.Conv2d(3, 6, 5) + self.pool = nn.MaxPool2d(2, 2) + self.conv2 = nn.Conv2d(6, 16, 5) + self.fc1 = nn.Linear(16 * 5 * 5, 120) + self.fc2 = nn.Linear(120, 84) + self.fc3 = nn.Linear(84, 10) + + def forward(self, x): + x = self.pool(F.relu(self.conv1(x))) + x = self.pool(F.relu(self.conv2(x))) + x = torch.flatten(x, 1) # flatten all dimensions except batch + x = F.relu(self.fc1(x)) + x = F.relu(self.fc2(x)) + x = self.fc3(x) + return x + + +def check_time(num_workers, batch_size, pin_memory, fast_loader): + transform = transforms.Compose( + transforms.ToImage(), + transforms.ToDtype(torch.float32, scale=True), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), + ) + + trainset = torchvision.datasets.CIFAR10(root=files.data_path(), train=True, download=False, transform=transform) + + loader_kwargs = dict(batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=pin_memory) + + if fast_loader: + trainloader = torch.utils.data.DataLoader(trainset, **loader_kwargs, persistent_workers=fast_loader) + else: + trainloader = torch.utils.data.DataLoader(trainset, **loader_kwargs, persistent_workers=fast_loader) + + net = Net() + net.to(DEVICE) + + criterion = nn.CrossEntropyLoss() + optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) + + time_per_batch = torch.zeros((2, len(trainloader))) + batch_time = time.time() + for cycle in range(2): + for idx, (images, labels) in enumerate(trainloader): + images, labels = images.to(DEVICE), labels.to(DEVICE) + + optimizer.zero_grad() + + outputs = net(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + time_per_batch[cycle, idx] = time.time() - batch_time + batch_time = time.time() + + avg_batch = torch.mean(time_per_batch[:, 1:]) + total_time = torch.sum(time_per_batch) + init_time = time_per_batch[0, 0] - avg_batch # estimate of prepare time + second_init = time_per_batch[1, 0] - avg_batch + return init_time, avg_batch, total_time, second_init + + +if __name__ == "__main__": + print("using device: ", DEVICE) + + # tests + num_workers = [0, 2, 4] + batch_size = [8, 1024] + pin_memory = [True, False] + use_fast_loader = [True, False] + + for nw in num_workers: + for bs in batch_size: + for pin in pin_memory: + for fast in use_fast_loader: + if nw == 0 and fast: + continue # persistent memory is irrelevant when num_workers=0 + + init, avg, total, second = check_time(nw, bs, pin, fast) + print(f"NW={nw}, BS={bs}, Pin={pin}, Persistent={fast}") + print(f"Dur={total:.3f}, PerBatch={avg:.2f}, Prep={init:.2f}, SecondPrep={second:.2f}") + print("") diff --git a/drafts/alignment_notes b/drafts/alignment_notes new file mode 160000 index 00000000..fea87596 --- /dev/null +++ b/drafts/alignment_notes @@ -0,0 +1 @@ +Subproject commit fea87596a5b84f587955c3af61fde1acc4359d5d diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 00000000..30ed34aa --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,34 @@ +# Development dependencies +-r requirements.txt + +# Testing +pytest>=7.0.0 +pytest-cov>=4.0.0 +pytest-xdist>=3.0.0 +pytest-benchmark>=4.0.0 + +# Code quality +black>=23.0.0 +isort>=5.12.0 +flake8>=6.0.0 +mypy>=1.5.0 +pre-commit>=3.0.0 + +# Security +bandit>=1.7.5 +safety>=2.3.0 +semgrep>=1.0.0 + +# Documentation +sphinx>=6.0.0 +sphinx-rtd-theme>=1.2.0 +sphinx-autodoc-typehints>=1.23.0 +linkchecker>=10.0.0 + +# Performance +memory-profiler>=0.60.0 +py-spy>=0.3.14 + +# Build +build>=0.10.0 +twine>=4.0.0 diff --git a/requirements.in b/requirements.in new file mode 100644 index 00000000..8e48ab8d --- /dev/null +++ b/requirements.in @@ -0,0 +1,37 @@ +# Core dependencies +torch>=1.12.0 +torchvision>=0.13.0 +transformers>=4.20.0 +datasets>=2.0.0 +accelerate>=0.20.0 + +# Scientific computing +numpy>=1.21.0 +scipy>=1.9.0 +scikit-learn>=1.1.0 +pandas>=1.4.0 + +# Visualization +matplotlib>=3.5.0 +seaborn>=0.11.0 +plotly>=5.10.0 + +# Configuration +pyyaml>=6.0 +omegaconf>=2.2.0 +hydra-core>=1.2.0 + +# Utilities +tqdm>=4.64.0 +rich>=12.0.0 +click>=8.1.0 +pathlib2>=2.3.0 + +# Data loading +pillow>=9.0.0 +opencv-python>=4.6.0 +h5py>=3.7.0 + +# HuggingFace +tokenizers>=0.13.0 +safetensors>=0.3.0