diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100644 index 0000000..ed30a72 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,65 @@ +name: Bug Report +description: Something is broken or producing incorrect results. +labels: ["bug"] +body: + - type: markdown + attributes: + value: | + Thanks for taking the time to report a bug. Please fill in as much detail as possible. + + - type: textarea + id: description + attributes: + label: Description + description: A clear and concise description of the bug. + validations: + required: true + + - type: textarea + id: reproduction + attributes: + label: Steps to reproduce + description: Minimal steps or code snippet to reproduce the issue. + placeholder: | + 1. Config used (paste relevant TOML section) + 2. Command run + 3. Error / unexpected output + validations: + required: true + + - type: textarea + id: expected + attributes: + label: Expected behaviour + description: What did you expect to happen? + validations: + required: true + + - type: textarea + id: environment + attributes: + label: Environment + description: | + Run `python -m ftnet.helper.collect_env` and paste the output here. + render: shell + validations: + required: true + + - type: dropdown + id: dataset + attributes: + label: Dataset + options: + - SODA + - MFN + - SCUT-Seg + - Cityscapes Thermal + - Other / N/A + validations: + required: false + + - type: textarea + id: additional + attributes: + label: Additional context + description: Logs, screenshots, or anything else that might help. diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 0000000..a65eafb --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,46 @@ +name: Feature Request +description: Suggest a new feature, dataset, or improvement. +labels: ["enhancement"] +body: + - type: markdown + attributes: + value: | + Thanks for suggesting an improvement! Please describe your idea clearly. + + - type: textarea + id: problem + attributes: + label: Problem / motivation + description: What problem does this solve? What is the current limitation? + validations: + required: true + + - type: textarea + id: solution + attributes: + label: Proposed solution + description: Describe the feature or change you'd like to see. + validations: + required: true + + - type: textarea + id: alternatives + attributes: + label: Alternatives considered + description: Any alternative approaches you have considered or tried. + + - type: dropdown + id: area + attributes: + label: Area + multiple: true + options: + - New dataset support + - Model architecture + - Training / optimisation + - Inference / deployment + - Documentation + - CI / tooling + - Other + validations: + required: true diff --git a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md new file mode 100644 index 0000000..2dd1039 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md @@ -0,0 +1,31 @@ +## Summary + + + +- + +## Motivation + + + +## Changes + + + +- + +## Testing + + + +- [ ] New unit tests added in `tests/` +- [ ] Existing tests pass (`poetry run pytest`) +- [ ] Pre-commit hooks pass (`poetry run pre-commit run --all-files`) + +## Checklist + +- [ ] PR targets `develop` (not `main`) +- [ ] Type hints added/updated for any new public functions +- [ ] Docstrings added/updated (Google style) +- [ ] Config changes reflected in `ftnet/cfg/cfg_dataclasses.py` with sensible defaults +- [ ] No `print()` statements — using `logger` instead diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..5608b85 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,97 @@ +name: CI + +on: + push: + branches: [develop, main] + pull_request: + branches: [develop, main] + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + name: Lint & Format + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install ruff + run: pip install ruff==0.4.4 + + - name: Run ruff lint + run: ruff check . + + - name: Run ruff format check + run: ruff format --check . + + type-check: + name: Type Check + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.10" + + - name: Install dependencies + run: pip install mypy==1.10.0 pydantic types-toml types-tqdm types-seaborn + + - name: Run mypy + run: mypy ftnet --ignore-missing-imports + + test: + name: Tests (Python ${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install Poetry + uses: snok/install-poetry@v1 + with: + virtualenvs-create: true + virtualenvs-in-project: true + + - name: Cache virtualenv + uses: actions/cache@v4 + with: + path: .venv + key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('poetry.lock') }} + restore-keys: | + venv-${{ runner.os }}-${{ matrix.python-version }}- + + - name: Install dependencies (CPU torch) + run: | + poetry config installer.max-workers 10 + poetry install --no-interaction --with dev -E cpu 2>/dev/null || \ + pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \ + poetry install --no-interaction --with dev --no-root && \ + pip install -e . + + - name: Run tests + run: poetry run pytest tests/ -v --cov=ftnet --cov-report=xml --cov-report=term-missing + + - name: Upload coverage + uses: codecov/codecov-action@v4 + if: matrix.python-version == '3.10' + with: + file: ./coverage.xml + fail_ci_if_error: false diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 21502c0..2e12cd6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -49,6 +49,13 @@ repos: #- mdformat-black - mdformat_frontmatter + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.10.0 + hooks: + - id: mypy + additional_dependencies: [pydantic, torch, types-toml, types-tqdm, types-seaborn] + args: [--ignore-missing-imports] + - repo: https://github.com/pre-commit/mirrors-prettier rev: v3.1.0 hooks: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..b16f949 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,132 @@ +# Contributing to FTNet + +Thank you for your interest in contributing! This document covers how to set up your development environment, the branch strategy, code standards, and how to submit changes. + +## Table of Contents + +- [Branch Strategy](#branch-strategy) +- [Development Setup](#development-setup) +- [Code Standards](#code-standards) +- [Running Tests](#running-tests) +- [Submitting a Pull Request](#submitting-a-pull-request) +- [Reporting Issues](#reporting-issues) + +--- + +## Branch Strategy + +| Branch | Purpose | +|--------|---------| +| `main` | Stable, published results matching the IEEE Access paper | +| `develop` | Active development — all PRs target this branch | + +**All pull requests must target `develop`, not `main`.** + +`main` is only updated via a release PR from `develop` when a new stable milestone is reached. + +--- + +## Development Setup + +### Prerequisites + +- Python 3.10+ +- [Poetry](https://python-poetry.org/docs/#installation) +- Git + +### Install + +```bash +git clone https://github.com/shreyaskamathkm/FTNet.git +cd FTNet +git checkout develop + +# Install all dependencies including dev tools +poetry install --with dev + +# Install pre-commit hooks +poetry run pre-commit install +``` + +### Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `FTNET_PRETRAINED_DIR` | `/pretrained_models` | Override path for pretrained backbone weights | + +--- + +## Code Standards + +This project uses [Ruff](https://docs.astral.sh/ruff/) for linting and formatting, and [mypy](https://mypy.readthedocs.io/) for static type checking. These are enforced automatically via pre-commit hooks. + +### Manual checks + +```bash +# Lint +poetry run ruff check . + +# Format +poetry run ruff format . + +# Type check +poetry run mypy ftnet --ignore-missing-imports +``` + +### Key conventions + +- **Type hints** are required on all public functions and methods. +- **Docstrings** follow Google style (Args / Returns sections). +- **Logging** — use `logger = logging.getLogger(__name__)` instead of `print()`. +- **Config changes** — all new parameters must be added to the appropriate Pydantic dataclass in `ftnet/cfg/cfg_dataclasses.py` with a sensible default. + +--- + +## Running Tests + +```bash +# Run all tests +poetry run pytest + +# Run with coverage report +poetry run pytest --cov=ftnet --cov-report=term-missing + +# Run a specific test file +poetry run pytest tests/test_loss.py -v +``` + +Tests live in the `tests/` directory and mirror the `ftnet/` package structure. + +--- + +## Submitting a Pull Request + +1. **Branch** off `develop`: + ```bash + git checkout develop + git pull + git checkout -b feat/your-feature-name + ``` + +2. **Make your changes** and write tests for any new functionality. + +3. **Ensure all checks pass** locally before opening a PR: + ```bash + poetry run pre-commit run --all-files + poetry run pytest + ``` + +4. **Open a PR** against `develop` (not `main`). Fill in the pull request template. + +5. A maintainer will review your PR. Please respond to review comments promptly. + +--- + +## Reporting Issues + +Use the GitHub issue templates: + +- **Bug report** — for incorrect behaviour, crashes, or wrong results. +- **Feature request** — for new dataset support, model variants, or tooling improvements. + +Please include your OS, Python version, GPU/CUDA version, and a minimal reproducible example where applicable. diff --git a/ftnet/cfg/cfg_dataclasses.py b/ftnet/cfg/cfg_dataclasses.py index 2998f64..31b929c 100644 --- a/ftnet/cfg/cfg_dataclasses.py +++ b/ftnet/cfg/cfg_dataclasses.py @@ -127,6 +127,7 @@ def validate_all_fields_at_the_same_time(cls, field_values): and not field_values.task.train_only ): raise ValueError( - "When train only mode is on, please update the dataset name to cityscapes_thermal_combine" + "Dataset 'cityscapes_thermal_combine' requires train_only=True. " + "Please set task.train_only = true in your config." ) return field_values diff --git a/ftnet/helper/model_helpers.py b/ftnet/helper/model_helpers.py index 56de658..ac4fcc0 100644 --- a/ftnet/helper/model_helpers.py +++ b/ftnet/helper/model_helpers.py @@ -73,11 +73,11 @@ def total_gradient(parameters: Any) -> float: def print_network(net: nn.Module) -> None: - """Print the network architecture and the total number of parameters. + """Log the network architecture and the total number of parameters. Args: - net (nn.Module): The network to print. + net (nn.Module): The network to log. """ num_params = sum(param.numel() for param in net.parameters()) - print(net) - print(f"Total number of parameters: {num_params / (1000**2)} M") + logger.info(net) + logger.info(f"Total number of parameters: {num_params / (1000**2):.2f} M") diff --git a/ftnet/models/encoder/resnetv1b.py b/ftnet/models/encoder/resnetv1b.py index e60ace2..ae6133c 100644 --- a/ftnet/models/encoder/resnetv1b.py +++ b/ftnet/models/encoder/resnetv1b.py @@ -15,7 +15,7 @@ from ...helper.model_helpers import check_mismatch from .splat import SplAtConv2d -root_pretrained_path = Path.cwd() / "pretrained_models" +root_pretrained_path = Path(os.environ.get("FTNET_PRETRAINED_DIR", Path(__file__).parents[3] / "pretrained_models")) root_pretrained_path.mkdir(parents=True, exist_ok=True) os.environ["TORCH_HOME"] = str(root_pretrained_path) diff --git a/pyproject.toml b/pyproject.toml index 47e23af..303c40e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,21 @@ priority = "explicit" [tool.poetry.group.dev.dependencies] pre-commit = "^3.7.0" ruff = "^0.4.4" +mypy = "^1.10.0" +pytest = "^8.2.0" +pytest-cov = "^5.0.0" + +[tool.mypy] +python_version = "3.10" +warn_return_any = true +warn_unused_configs = true +ignore_missing_imports = true +disallow_untyped_defs = true +disallow_incomplete_defs = true + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "--cov=ftnet --cov-report=term-missing --cov-report=xml" [build-system] requires = ["poetry-core"] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_cfg_dataclasses.py b/tests/test_cfg_dataclasses.py new file mode 100644 index 0000000..a63df11 --- /dev/null +++ b/tests/test_cfg_dataclasses.py @@ -0,0 +1,118 @@ +"""Tests for FTNet configuration dataclasses.""" + +import pytest +from pydantic import ValidationError + +from ftnet.cfg.cfg_dataclasses import ( + CheckpointLogArgs, + ComputeArgs, + DataLoaderArgs, + FTNetArgs, + ModelArgs, + OptimizerArgs, + SchedulerArgs, + TaskArgs, + TrainingHyperParamsArgs, + WandBArgs, +) + + +class TestTaskArgs: + def test_defaults(self): + args = TaskArgs() + assert args.mode == "train" + assert args.train_only is False + assert args.debug is False + + def test_valid_modes(self): + for mode in ("train", "test", "infer"): + assert TaskArgs(mode=mode).mode == mode + + def test_invalid_mode(self): + with pytest.raises(ValidationError): + TaskArgs(mode="invalid") + + +class TestModelArgs: + def test_defaults(self): + args = ModelArgs() + assert args.name == "ftnet" + assert args.backbone == "resnext50_32x4d" + assert args.no_of_filters == 128 + assert args.edge_extracts == [3] + assert args.num_blocks == 2 + + def test_custom_values(self): + args = ModelArgs(backbone="resnet50", no_of_filters=256, num_blocks=4) + assert args.backbone == "resnet50" + assert args.no_of_filters == 256 + assert args.num_blocks == 4 + + +class TestDataLoaderArgs: + def test_defaults(self): + args = DataLoaderArgs() + assert args.name == "soda" + + def test_valid_dataset_names(self): + valid = ["cityscapes_thermal_combine", "cityscapes_thermal_split", "soda", "mfn", "scutseg"] + for name in valid: + assert DataLoaderArgs(name=name).name == name + + def test_invalid_dataset_name(self): + with pytest.raises(ValidationError): + DataLoaderArgs(name="nonexistent_dataset") + + +class TestFTNetArgs: + def test_defaults(self): + args = FTNetArgs() + assert isinstance(args.task, TaskArgs) + assert isinstance(args.trainer, TrainingHyperParamsArgs) + assert isinstance(args.model, ModelArgs) + assert isinstance(args.dataset, DataLoaderArgs) + assert isinstance(args.wandb, WandBArgs) + assert isinstance(args.optimizer, OptimizerArgs) + assert isinstance(args.scheduler, SchedulerArgs) + assert isinstance(args.checkpoint, CheckpointLogArgs) + assert isinstance(args.compute, ComputeArgs) + + def test_cityscapes_combine_requires_train_only(self): + """cityscapes_thermal_combine must have train_only=True.""" + with pytest.raises(ValidationError, match="requires train_only=True"): + FTNetArgs( + dataset=DataLoaderArgs(name="cityscapes_thermal_combine"), + task=TaskArgs(train_only=False), + ) + + def test_cityscapes_combine_with_train_only_passes(self): + """cityscapes_thermal_combine with train_only=True should be valid.""" + args = FTNetArgs( + dataset=DataLoaderArgs(name="cityscapes_thermal_combine"), + task=TaskArgs(train_only=True), + ) + assert args.dataset.name == "cityscapes_thermal_combine" + assert args.task.train_only is True + + def test_other_datasets_without_train_only_passes(self): + """Non-cityscapes_thermal_combine datasets should not require train_only.""" + for name in ("soda", "mfn", "scutseg"): + args = FTNetArgs( + dataset=DataLoaderArgs(name=name), + task=TaskArgs(train_only=False), + ) + assert args.dataset.name == name + + def test_from_config_with_dict(self): + config = {"task": {"mode": "train"}, "dataset": {"name": "soda"}} + args = FTNetArgs.from_config(config) + assert args.task.mode == "train" + assert args.dataset.name == "soda" + + def test_from_config_rejects_extra_keys(self): + with pytest.raises(ValueError, match="Unexpected keys"): + FTNetArgs.from_config({"nonexistent_section": {}}) + + def test_from_config_invalid_type(self): + with pytest.raises(ValueError, match="Invalid input type"): + FTNetArgs.from_config(object()) # type: ignore[arg-type] diff --git a/tests/test_data_attributes.py b/tests/test_data_attributes.py new file mode 100644 index 0000000..19291a3 --- /dev/null +++ b/tests/test_data_attributes.py @@ -0,0 +1,31 @@ +"""Tests for dataloader data attribute dataclasses.""" + +from ftnet.core.dataloaders.data_attributes import ImageSize, ImageSizes + + +class TestImageSize: + def test_stores_dimensions(self): + size = ImageSize(width=640, height=480) + assert size.width == 640 + assert size.height == 480 + + def test_equality(self): + assert ImageSize(100, 200) == ImageSize(100, 200) + assert ImageSize(100, 200) != ImageSize(200, 100) + + +class TestImageSizes: + def test_stores_base_and_crop(self): + base = ImageSize(width=640, height=480) + crop = ImageSize(width=512, height=512) + sizes = ImageSizes(base_size=base, crop_size=crop) + assert sizes.base_size == base + assert sizes.crop_size == crop + + def test_fields_are_accessible(self): + sizes = ImageSizes( + base_size=ImageSize(320, 240), + crop_size=ImageSize(256, 256), + ) + assert sizes.base_size.width == 320 + assert sizes.crop_size.height == 256 diff --git a/tests/test_loss.py b/tests/test_loss.py new file mode 100644 index 0000000..dc5ac8f --- /dev/null +++ b/tests/test_loss.py @@ -0,0 +1,84 @@ +"""Tests for segmentation loss functions.""" + +import pytest +import torch + +from ftnet.core.loss.segmentation_loss import EdgeNetLoss, MixSoftmaxCrossEntropyLoss + + +class TestMixSoftmaxCrossEntropyLoss: + def test_basic_forward(self): + loss_fn = MixSoftmaxCrossEntropyLoss() + pred = torch.randn(2, 5, 8, 8) + target = torch.randint(0, 5, (2, 8, 8)) + loss = loss_fn(pred, target) + assert loss.ndim == 0 + assert loss.item() > 0 + + def test_ignore_index(self): + loss_fn = MixSoftmaxCrossEntropyLoss(ignore_index=255) + pred = torch.randn(2, 5, 8, 8) + target = torch.full((2, 8, 8), 255, dtype=torch.long) + loss = loss_fn(pred, target) + # All targets ignored — loss should be 0 or nan-safe + assert not torch.isnan(loss) + + def test_aux_loss_reduces_with_correct_pred(self): + loss_fn = MixSoftmaxCrossEntropyLoss(aux=True, aux_weight=0.4) + # Provide a tuple of (main_pred, aux_pred) and target + main_pred = torch.randn(2, 5, 8, 8) + aux_pred = torch.randn(2, 5, 8, 8) + target = torch.randint(0, 5, (2, 8, 8)) + loss = loss_fn((main_pred, aux_pred), target) + assert loss.ndim == 0 + assert loss.item() > 0 + + def test_aux_forward_combines_losses(self): + loss_fn = MixSoftmaxCrossEntropyLoss(aux=True, aux_weight=0.5) + pred_main = torch.randn(2, 5, 8, 8) + pred_aux = torch.randn(2, 5, 8, 8) + target = torch.randint(0, 5, (2, 8, 8)) + loss = loss_fn(pred_main, pred_aux, target) + assert torch.isfinite(loss) + + +class TestEdgeNetLoss: + def _make_inputs(self, batch=2, num_classes=5, h=8, w=8): + pred_class = torch.randn(batch, num_classes, h, w) + pred_edge = torch.randn(batch, 1, h, w) + target_class = torch.randint(0, num_classes, (batch, h, w)) + target_edge = torch.randint(0, 2, (batch, h, w)).float() + return (pred_class, pred_edge), (target_class, target_edge) + + def test_forward_returns_scalar(self): + loss_fn = EdgeNetLoss(loss_weight=20) + preds, targets = self._make_inputs() + loss = loss_fn(preds, targets) + assert loss.ndim == 0 + assert torch.isfinite(loss) + + def test_loss_weight_scales_edge_loss(self): + torch.manual_seed(42) + preds, targets = self._make_inputs() + + loss_low = EdgeNetLoss(loss_weight=1)(preds, targets) + loss_high = EdgeNetLoss(loss_weight=50)(preds, targets) + # Higher weight should produce a higher total loss + assert loss_high.item() > loss_low.item() + + def test_requires_tuple_input(self): + loss_fn = EdgeNetLoss() + pred = torch.randn(2, 5, 8, 8) + target = torch.randint(0, 5, (2, 8, 8)) + with pytest.raises((ValueError, TypeError)): + loss_fn(pred, target) + + def test_auto_weight_bce_3d_target(self): + """Edge target with 3 dims (no channel) should be handled correctly.""" + loss_fn = EdgeNetLoss(loss_weight=1) + pred_class = torch.randn(2, 5, 8, 8) + pred_edge = torch.randn(2, 1, 8, 8) + target_class = torch.randint(0, 5, (2, 8, 8)) + target_edge = torch.randint(0, 2, (2, 8, 8)).float() # 3D: no channel dim + loss = loss_fn((pred_class, pred_edge), (target_class, target_edge)) + assert torch.isfinite(loss) diff --git a/tests/test_model_helpers.py b/tests/test_model_helpers.py new file mode 100644 index 0000000..e435720 --- /dev/null +++ b/tests/test_model_helpers.py @@ -0,0 +1,82 @@ +"""Tests for model helper utilities.""" + +from collections import OrderedDict +from pathlib import Path + +import torch +import torch.nn as nn + +from ftnet.helper.model_helpers import check_mismatch, save_model_summary, total_gradient + + +class TestCheckMismatch: + def _make_state_dict(self, keys_shapes: dict) -> dict: + return {k: torch.zeros(shape) for k, shape in keys_shapes.items()} + + def test_matching_keys_and_shapes(self): + model_dict = self._make_state_dict({"layer.weight": (4, 4), "layer.bias": (4,)}) + # pretrained keys have a 6-char prefix (e.g. "model.") + pretrained_dict = self._make_state_dict( + {"model.layer.weight": (4, 4), "model.layer.bias": (4,)} + ) + result = check_mismatch(model_dict, pretrained_dict) + assert set(result.keys()) == {"layer.weight", "layer.bias"} + + def test_shape_mismatch_skips_key(self): + model_dict = self._make_state_dict({"layer.weight": (4, 4)}) + pretrained_dict = self._make_state_dict({"model.layer.weight": (8, 8)}) + result = check_mismatch(model_dict, pretrained_dict) + assert "layer.weight" not in result + + def test_missing_key_in_model_skipped(self): + model_dict = self._make_state_dict({"layer.weight": (4, 4)}) + pretrained_dict = self._make_state_dict({"model.other.weight": (4, 4)}) + result = check_mismatch(model_dict, pretrained_dict) + assert len(result) == 0 + + def test_returns_ordered_dict(self): + model_dict = self._make_state_dict({"layer.weight": (4, 4)}) + pretrained_dict = self._make_state_dict({"model.layer.weight": (4, 4)}) + result = check_mismatch(model_dict, pretrained_dict) + assert isinstance(result, OrderedDict) + + def test_empty_pretrained(self): + model_dict = self._make_state_dict({"layer.weight": (4, 4)}) + result = check_mismatch(model_dict, {}) + assert len(result) == 0 + + +class TestSaveModelSummary: + def test_creates_file(self, tmp_path: Path): + model = nn.Linear(4, 2) + save_model_summary(model, tmp_path) + summary_file = tmp_path / "model.txt" + assert summary_file.exists() + + def test_file_contains_param_count(self, tmp_path: Path): + model = nn.Linear(4, 2) + save_model_summary(model, tmp_path) + content = (tmp_path / "model.txt").read_text() + assert "Total number of parameters" in content + + def test_file_contains_model_repr(self, tmp_path: Path): + model = nn.Linear(4, 2) + save_model_summary(model, tmp_path) + content = (tmp_path / "model.txt").read_text() + assert "Linear" in content + + +class TestTotalGradient: + def test_with_gradients(self): + model = nn.Linear(4, 2) + output = model(torch.randn(3, 4)) + output.sum().backward() + norm = total_gradient(model.parameters()) + assert isinstance(norm, float) + assert norm > 0 + + def test_no_gradients_returns_zero(self): + model = nn.Linear(4, 2) + # No backward pass — no gradients + norm = total_gradient(model.parameters()) + assert norm == 0.0 diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py new file mode 100644 index 0000000..7479c69 --- /dev/null +++ b/tests/test_schedulers.py @@ -0,0 +1,99 @@ +"""Tests for learning rate schedulers.""" + +import pytest +import torch +import torch.nn as nn +from torch.optim import SGD + +from ftnet.core.schedulers.warmup_lr import WarmupMultiStepLR, WarmupPolyLR + + +def _make_optimizer(lr: float = 0.1) -> SGD: + model = nn.Linear(2, 2) + return SGD(model.parameters(), lr=lr) + + +class TestWarmupMultiStepLR: + def test_invalid_milestones_raises(self): + opt = _make_optimizer() + with pytest.raises(ValueError, match="increasing"): + WarmupMultiStepLR(opt, milestones=[10, 5]) + + def test_invalid_warmup_method_raises(self): + opt = _make_optimizer() + with pytest.raises(ValueError, match="warmup_method"): + WarmupMultiStepLR(opt, milestones=[10], warmup_method="cosine") + + def test_lr_decreases_at_milestone(self): + opt = _make_optimizer(lr=0.1) + scheduler = WarmupMultiStepLR(opt, milestones=[2], gamma=0.1, warmup_iters=0) + lrs_before = scheduler.get_last_lr() if hasattr(scheduler, "get_last_lr") else [0.1] + # Step past the milestone + for _ in range(3): + opt.step() + scheduler.step() + lr_after = scheduler.get_last_lr()[0] + assert lr_after < 0.1 + + def test_warmup_linear_increases_lr(self): + opt = _make_optimizer(lr=0.1) + warmup_iters = 5 + scheduler = WarmupMultiStepLR( + opt, + milestones=[100], + warmup_factor=1.0 / 3, + warmup_iters=warmup_iters, + warmup_method="linear", + ) + lrs = [] + for _ in range(warmup_iters): + opt.step() + scheduler.step() + lrs.append(scheduler.get_last_lr()[0]) + # LR should increase monotonically during warmup + assert lrs == sorted(lrs) + + def test_warmup_constant(self): + opt = _make_optimizer(lr=0.1) + scheduler = WarmupMultiStepLR( + opt, + milestones=[100], + warmup_factor=1.0 / 3, + warmup_iters=3, + warmup_method="constant", + ) + opt.step() + scheduler.step() + lr = scheduler.get_last_lr()[0] + assert lr < 0.1 # warmup_factor < 1 means LR is reduced during warmup + + +class TestWarmupPolyLR: + def test_invalid_warmup_method_raises(self): + opt = _make_optimizer() + with pytest.raises(ValueError, match="warmup_method"): + WarmupPolyLR(opt, steps_per_epoch=10, epochs=10, warmup_method="cosine") + + def test_lr_decays_over_time(self): + opt = _make_optimizer(lr=0.1) + scheduler = WarmupPolyLR(opt, steps_per_epoch=10, epochs=5, warmup_iters=0) + lrs = [] + for _ in range(40): + opt.step() + scheduler.step() + lrs.append(scheduler.get_last_lr()[0]) + # LR should generally decrease (poly decay) + assert lrs[0] >= lrs[-1] + + def test_reaches_target_lr(self): + target_lr = 1e-6 + opt = _make_optimizer(lr=0.1) + total_steps = 10 + scheduler = WarmupPolyLR( + opt, steps_per_epoch=total_steps, epochs=1, target_lr=target_lr, warmup_iters=0 + ) + for _ in range(total_steps): + opt.step() + scheduler.step() + final_lr = scheduler.get_last_lr()[0] + assert final_lr >= target_lr