Improve training: learnable temperature, per-group LR, warmup, augmentations

Loss:
- Learnable temperature (CLIP-style logit_scale) with clamp [0.01, 0.5]
- Replaces fixed cosine schedule (still available via --no-learnable-temp)
- Default tau_init=0.07

Optimizer:
- Per-group LR: projections 1e-4, text encoder 1e-5 (10x lower)
- Learnable temperature included in projection param group

Scheduler:
- Linear warmup (2 epochs default) + cosine annealing
- Per-step scheduling (not per-epoch)

Augmentations (separate drone/satellite):
- Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15), ColorJitter,
  RandomGrayscale(0.05), GaussianBlur
- Satellite: RandomResizedCrop(0.7-1.0), HFlip, ColorJitter, RandomGrayscale
- Eval: clean Resize+CenterCrop (no augmentation)

Dataset: supports separate drone_transform/sat_transform args

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 18:07:17 +03:00
parent 6ad9c4d149
commit 998d52cb57
4 changed files with 210 additions and 41 deletions

View File

@@ -543,15 +543,58 @@ class AsymmetricEncoder(nn.Module):
# Image preprocessing (DINOv3: 256x256, ImageNet normalization)
# ---------------------------------------------------------------------------
_IMAGENET_MEAN = [0.485, 0.456, 0.406]
_IMAGENET_STD = [0.229, 0.224, 0.225]
def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
"""Build image transform for DINOv3 input."""
"""Build eval/inference image transform for DINOv3 input."""
from torchvision import transforms
return transforms.Compose([
transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
])
def get_drone_train_transform(image_size: int = 256) -> torch.nn.Module:
"""Build training augmentation for drone images.
Includes: RandomResizedCrop, HFlip, rotation, color jitter,
grayscale, Gaussian blur.
"""
from torchvision import transforms
return transforms.Compose([
transforms.RandomResizedCrop(
image_size, scale=(0.7, 1.0),
interpolation=transforms.InterpolationMode.BICUBIC,
),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomRotation(degrees=15),
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
transforms.RandomGrayscale(p=0.05),
transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
transforms.ToTensor(),
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
])
def get_satellite_train_transform(image_size: int = 256) -> torch.nn.Module:
"""Build training augmentation for satellite images.
Lighter than drone: no rotation or blur (satellite is nadir/consistent).
Includes: RandomResizedCrop, HFlip, color jitter, grayscale.
"""
from torchvision import transforms
return transforms.Compose([
transforms.RandomResizedCrop(
image_size, scale=(0.7, 1.0),
interpolation=transforms.InterpolationMode.BICUBIC,
),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
transforms.RandomGrayscale(p=0.05),
transforms.ToTensor(),
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
])