diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index 85e66e0..59b2fb4 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -2,18 +2,19 @@ name: Publish ApDepth Docker Image on: push: - branches: - - main tags: - "v*" - workflow_dispatch: + +concurrency: + group: docker-release-${{ github.ref }} + cancel-in-progress: false env: REGISTRY: ghcr.io IMAGE_NAME: haruko386/apdepth permissions: - contents: read + contents: write packages: write jobs: @@ -48,4 +49,19 @@ jobs: file: ./Dockerfile push: true tags: ${{ steps.meta.outputs.tags }} - labels: ${{ steps.meta.outputs.labels }} \ No newline at end of file + labels: ${{ steps.meta.outputs.labels }} + + - name: Create GitHub Release + env: + GH_TOKEN: ${{ github.token }} + TAG_NAME: ${{ github.ref_name }} + run: | + if gh release view "${TAG_NAME}" --repo "${GITHUB_REPOSITORY}" >/dev/null 2>&1; then + echo "Release ${TAG_NAME} already exists; skipping creation." + else + gh release create "${TAG_NAME}" \ + --repo "${GITHUB_REPOSITORY}" \ + --title "${TAG_NAME}" \ + --generate-notes \ + --verify-tag + fi diff --git a/README.md b/README.md index 87dd845..8a773e5 100755 --- a/README.md +++ b/README.md @@ -238,6 +238,14 @@ To train the main stage directly from the base model without Stage 1 initializat python train.py --config config/train_apdepth.yaml --no_wandb ``` +The transition to latent frequency loss is controlled by +`latent_freq_loss.gradual_transition` in `config/train_apdepth.yaml`. When `true`, +the reconstruction loss weight decreases linearly from 1 to 0 between iteration +20,000 and `max_iter`, while the frequency-loss weight increases from 0 to 1. +When `false` or omitted, training preserves the original hard switch immediately +after iteration 20,000. TensorBoard records the active frequency weight as +`train/freq_loss_weight`. + Resume from a checkpoint, e.g. ```bash diff --git a/config/train_apdepth.yaml b/config/train_apdepth.yaml index feeed27..20fc515 100644 --- a/config/train_apdepth.yaml +++ b/config/train_apdepth.yaml @@ -65,6 +65,9 @@ latent_freq_loss: loss_type: 'l1' high_pass_weight: 1.0 lambda: 0.5 + # true: linearly blend reconstruction -> frequency loss from iter 20k to max_iter + # false: preserve the original hard switch immediately after iter 20k + gradual_transition: true lr: 3.0e-05 # Stage 2 uses the original objective. Enable known-far supervision only in diff --git a/src/trainer/apdepth_trainer.py b/src/trainer/apdepth_trainer.py index a84eba5..82e791d 100644 --- a/src/trainer/apdepth_trainer.py +++ b/src/trainer/apdepth_trainer.py @@ -54,6 +54,8 @@ from src.util.seeding import generate_seed_sequence class ApDepthTrainer: + FREQUENCY_LOSS_START_ITER = 20000 + def __init__( self, cfg: OmegaConf, @@ -114,6 +116,11 @@ def __init__( self.loss = get_loss(loss_name=self.cfg.loss.name, **self.cfg.loss.kwargs) self.l1_loss = get_loss(loss_name=self.cfg.l1_loss.name, **self.cfg.l1_loss.kwargs) self.latent_freq_loss = get_loss(loss_name=self.cfg.latent_freq_loss.name, ** self.cfg.latent_freq_loss.kwargs) + self.gradual_loss_transition = self.cfg.latent_freq_loss.get( + "gradual_transition", False + ) + if not isinstance(self.gradual_loss_transition, bool): + raise TypeError("latent_freq_loss.gradual_transition must be a boolean") far_cfg = self.cfg.get("far_depth_supervision", {}) self.far_enabled = bool(far_cfg.get("enabled", False)) @@ -129,7 +136,12 @@ def __init__( # Eval metrics self.metric_funcs = [getattr(metric, _met) for _met in cfg.eval.eval_metrics] self.train_metrics = MetricTracker( - "loss", "latent_freq_loss", "far_loss", "far_pixel_ratio", "far_latent_ratio" + "loss", + "latent_freq_loss", + "freq_loss_weight", + "far_loss", + "far_pixel_ratio", + "far_latent_ratio", ) self.val_metrics = MetricTracker(*[m.__name__ for m in self.metric_funcs]) @@ -259,12 +271,13 @@ def train(self, t_end=None): torch.cuda.empty_cache() - loss, freq_loss, far_loss = self._compute_depth_loss( + loss, freq_loss, freq_loss_weight, far_loss = self._compute_depth_loss( depth_pred, gt_depth_latent, depth, depth_gt_for_loss, valid_mask_down, valid_mask_for_latent, far_mask_down, far_mask, ) self.train_metrics.update("loss", loss.item()) self.train_metrics.update("latent_freq_loss", freq_loss.item()) + self.train_metrics.update("freq_loss_weight", freq_loss_weight) self.train_metrics.update("far_loss", far_loss.item()) self.train_metrics.update("far_pixel_ratio", far_mask.float().mean().item()) self.train_metrics.update("far_latent_ratio", far_mask_down.float().mean().item()) @@ -342,13 +355,29 @@ def _compute_depth_loss( valid_latent, valid_pixel, far_latent, far_pixel, ): # Preserve the original valid-GT objective and its normalization. + freq_loss_weight = self._get_frequency_loss_weight() + reconstruction_loss_weight = 1.0 - freq_loss_weight + + reconstruction_loss = pred_latent.new_zeros(()) freq_loss = pred_latent.new_zeros(()) - if self.effective_iter <= 20000: - loss = masked_loss(self.loss, pred_latent.float(), target_latent.float(), valid_latent) - loss = loss + masked_loss(self.l1_loss, pred_depth.float(), target_depth.float(), valid_pixel) - else: - freq_loss = self.latent_freq_loss(pred_latent.float(), target_latent.float(), valid_latent) - loss = self.cfg.latent_freq_loss.get("lambda", 1.0) * freq_loss + if reconstruction_loss_weight > 0: + reconstruction_loss = masked_loss( + self.loss, pred_latent.float(), target_latent.float(), valid_latent + ) + reconstruction_loss = reconstruction_loss + masked_loss( + self.l1_loss, pred_depth.float(), target_depth.float(), valid_pixel + ) + if freq_loss_weight > 0: + freq_loss = self.latent_freq_loss( + pred_latent.float(), target_latent.float(), valid_latent + ) + + loss = reconstruction_loss_weight * reconstruction_loss + loss = loss + ( + freq_loss_weight + * self.cfg.latent_freq_loss.get("lambda", 1.0) + * freq_loss + ) far_loss = pred_latent.new_zeros(()) if self.far_enabled: @@ -357,7 +386,20 @@ def _compute_depth_loss( far_loss = masked_loss(self.loss, pred_latent.float(), target_latent.float(), far_latent) far_loss = far_loss + masked_loss(self.l1_loss, pred_depth.float(), target_depth.float(), far_pixel) loss = loss + self.far_weight * far_loss - return loss, freq_loss, far_loss + return loss, freq_loss, freq_loss_weight, far_loss + + def _get_frequency_loss_weight(self): + if not self.gradual_loss_transition: + return float(self.effective_iter > self.FREQUENCY_LOSS_START_ITER) + + transition_length = self.max_iter - self.FREQUENCY_LOSS_START_ITER + if transition_length <= 0: + return float(self.effective_iter > self.FREQUENCY_LOSS_START_ITER) + + progress = ( + self.effective_iter - self.FREQUENCY_LOSS_START_ITER + ) / transition_length + return min(max(progress, 0.0), 1.0) def encode_depth(self, depth_in): # stack depth into 3-channel