Add integration example

examples/integration/README.md

@@ -0,0 +1,26 @@
+# Example of Black Box Optimization on ABCI 3.0
+
+This is an example of performing black-box optimization of the learning rate for a ResNet50 model on the MNIST dataset.
+
+## Getting started
+
+In an environment where aiaccel is installed, additionally install torchvision.
+
+```bash
+pip install torchvision
+```
+
+
+Run the following command to perform black-box optimization.
+PATH_TO_ENV should be changed to the path of the environment prepared above.
+
+```bash
+python -m aiaccel.hpo.apps.optimize "python -m aiaccel.jobs.cli.abci3 gpu --command_prefix 'cd \$PBS_O_WORKDIR && module load cuda/12.6/12.6.1 && module load python/3.13/3.13.2 && source PATH_TO_ENV/bin/activate &&' jobs/{job_name}.log -- python -m aiaccel.torch.apps.train resnet50/config.yaml task.optimizer_config.optimizer_generator.lr={lr} trainer.logger.name=lr_{lr} out_filename={out_filename}" --config config.yaml
+```
+
+## Detailed Descriptions
+
+The target function for optimization using aiaccel.hpo.app.optimize is objective_integration.main.
+Within objective_integration.main, aiaccel.torch.app.train is called, and the learning rate is returned.
+
+Detailed descriptions of torch and optimize are available on the [aiaccel document(torch)](https://aistairc.github.io/aiaccel/user_guide/torch.html) [aiaccel document(optimize)](https://aistairc.github.io/aiaccel/user_guide/hpo.html)

examples/integration/common_config.yaml

@@ -0,0 +1,61 @@
+trainer:
+  max_epochs: 10
+
+  callbacks:
+    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
+      filename: "{epoch:04d}"
+      save_last: True
+      save_top_k: -1
+    - _target_: torchvision_task_integration.SaveValLossCallback
+      output_path: ${out_filename}
+
+
+datamodule:
+  _target_: aiaccel.torch.lightning.datamodules.single_datamodule.SingleDataModule
+
+  train_dataset_fn:
+    _partial_: true
+    _target_: torchvision.datasets.MNIST
+    train: True
+
+  val_dataset_fn:
+    _partial_: true
+    _target_: torchvision.datasets.MNIST
+    train: False
+
+  common_args:
+    root: "./dataset"
+    download: True
+    transform: ${transform}
+
+  batch_size: 128
+  use_scatter: False
+
+
+transform:
+  _target_: torchvision.transforms.Compose
+  transforms:
+    - _target_: torchvision.transforms.Resize
+      size: [256, 256]
+    - _target_: torchvision.transforms.Grayscale
+      num_output_channels: 3
+    - _target_: torchvision.transforms.ToTensor
+    - _target_: torchvision.transforms.Normalize
+      mean: [0.5]
+      std: [0.5]
+
+task:
+  _target_: torchvision_task_integration.ImageClassificationTask
+  num_classes: 10
+
+  model:
+    _target_: torchvision.models.resnet50
+    weights:
+      _target_: hydra.utils.get_object
+      path: torchvision.models.ResNet50_Weights.DEFAULT
+
+  optimizer_config:
+    _target_: aiaccel.torch.lightning.OptimizerConfig
+    optimizer_generator:
+      _partial_: True
+      _target_: torch.optim.Adam

examples/integration/config.yaml

@@ -0,0 +1,12 @@
+params:
+  _convert_: partial
+  _target_: aiaccel.hpo.apps.optimize.HparamsManager
+  lr:
+    _target_: aiaccel.hpo.optuna.suggest_wrapper.SuggestFloat
+    name: lr
+    low: 1.e-6
+    high: 1.e-2
+    log: true
+
+n_trials: 30
+n_max_jobs: 4

examples/integration/resnet50/config.yaml

@@ -0,0 +1,3 @@
+_base_:
+  - ${base_config_path}/train_base.yaml
+  - ../common_config.yaml

examples/integration/torchvision_task_integration.py

@@ -0,0 +1,75 @@
+import json
+
+import torch
+from torch import nn
+from torch.nn import functional as func
+
+import lightning.pytorch as pl
+
+from torchmetrics.classification import MulticlassAccuracy
+
+from aiaccel.torch.lightning import OptimizerConfig, OptimizerLightningModule
+
+
+class ImageClassificationTask(OptimizerLightningModule):
+    def __init__(self, model: nn.Module, optimizer_config: OptimizerConfig, num_classes: int = 10):
+        super().__init__(optimizer_config)
+
+        self.model = model
+        if hasattr(self.model.fc, "in_features") and isinstance(self.model.fc.in_features, int):
+            self.model.fc = nn.Linear(self.model.fc.in_features, num_classes)
+
+        self.train_accuracy = MulticlassAccuracy(num_classes=num_classes)
+        self.val_accuracy = MulticlassAccuracy(num_classes=num_classes)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)  # type: ignore
+
+    def training_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        x, y = batch
+
+        logits = self(x)
+
+        loss = func.cross_entropy(logits, y)
+
+        acc = self.train_accuracy(logits, y)
+        self.log_dict(
+            {
+                "training/loss": loss,
+                "training/acc": acc,
+            },
+            prog_bar=True,
+        )
+
+        return loss
+
+    def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
+        x, y = batch
+
+        logits = self(x)
+
+        loss = func.cross_entropy(logits, y)
+
+        acc = self.val_accuracy(logits, y)
+        self.log_dict(
+            {
+                "validation/loss": loss,
+                "validation/acc": acc,
+            },
+            prog_bar=True,
+        )
+
+
+class SaveValLossCallback(pl.Callback):
+    def __init__(self, output_path: str) -> None:
+        super().__init__()
+        self.output_path = output_path
+
+    def on_fit_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule) -> None:
+        val_loss = trainer.callback_metrics.get("validation/loss")
+        if val_loss is not None:
+            val_loss_value = val_loss.item()
+            with open(self.output_path, "w") as f:
+                json.dump(val_loss_value, f)
+        else:
+            print("Warning: 'validation/loss' not found in callback_metrics.")