64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import os
|
|
import tempfile
|
|
from datetime import datetime
|
|
from typing import Optional
|
|
|
|
import pandas as pd
|
|
import torch
|
|
from optuna.distributions import CategoricalDistribution, LogUniformDistribution
|
|
from torchmetrics import Accuracy
|
|
|
|
import lightning as L
|
|
from lightning.app.components import TracerPythonScript
|
|
|
|
|
|
class ObjectiveWork(TracerPythonScript):
|
|
def __init__(self, script_path: str, data_dir: str, cloud_compute: Optional[L.CloudCompute]):
|
|
timestamp = datetime.now().strftime("%H:%M:%S")
|
|
tmpdir = tempfile.TemporaryDirectory().name
|
|
submission_path = os.path.join(tmpdir, f"{timestamp}.csv")
|
|
best_model_path = os.path.join(tmpdir, f"{timestamp}.model.pt")
|
|
super().__init__(
|
|
script_path,
|
|
script_args=[
|
|
f"--train_data_path={data_dir}/train",
|
|
f"--test_data_path={data_dir}/test",
|
|
f"--submission_path={submission_path}",
|
|
f"--best_model_path={best_model_path}",
|
|
],
|
|
cloud_compute=cloud_compute,
|
|
)
|
|
self.data_dir = data_dir
|
|
self.best_model_path = best_model_path
|
|
self.submission_path = submission_path
|
|
self.metric = None
|
|
self.trial_id = None
|
|
self.metric = None
|
|
self.params = None
|
|
self.has_told_study = False
|
|
|
|
def run(self, trial_id: int, **params):
|
|
self.trial_id = trial_id
|
|
self.params = params
|
|
self.script_args.extend([f"--{k}={v}" for k, v in params.items()])
|
|
super().run()
|
|
self.compute_metric()
|
|
|
|
def _to_labels(self, path: str):
|
|
return torch.from_numpy(pd.read_csv(path).label.values)
|
|
|
|
def compute_metric(self):
|
|
self.metric = -1 * float(
|
|
Accuracy()(
|
|
self._to_labels(self.submission_path),
|
|
self._to_labels(f"{self.data_dir}/ground_truth.csv"),
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def distributions():
|
|
return {
|
|
"backbone": CategoricalDistribution(["resnet18", "resnet34"]),
|
|
"learning_rate": LogUniformDistribution(0.0001, 0.1),
|
|
}
|