Suppress spurious lr_scheduler.step() warning from PyTorch

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-21 18:24:54 +03:00
parent 517b87d3d8
commit fa32b2e67f

View File

@@ -11,6 +11,7 @@ import json
import logging import logging
import math import math
import time import time
import warnings
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
@@ -317,10 +318,13 @@ def train(cfg: TrainConfigGTAUAV) -> None:
steps_per_epoch = len(train_loader) steps_per_epoch = len(train_loader)
total_steps = cfg.epochs * steps_per_epoch total_steps = cfg.epochs * steps_per_epoch
warmup_steps = cfg.warmup_epochs * steps_per_epoch warmup_steps = cfg.warmup_epochs * steps_per_epoch
scheduler = LambdaLR( with warnings.catch_warnings():
optimizer, warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps), scheduler = LambdaLR(
) optimizer,
lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
last_epoch=-1,
)
scaler = GradScaler(enabled=cfg.use_amp) scaler = GradScaler(enabled=cfg.use_amp)
# Restore optimizer/scheduler/loss state on resume. # Restore optimizer/scheduler/loss state on resume.
@@ -385,7 +389,9 @@ def train(cfg: TrainConfigGTAUAV) -> None:
) )
scaler.step(optimizer) scaler.step(optimizer)
scaler.update() scaler.update()
scheduler.step() with warnings.catch_warnings():
warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*")
scheduler.step()
for key, val in loss_dict.items(): for key, val in loss_dict.items():
agg[key] = agg.get(key, 0.0) + float(val.item()) agg[key] = agg.get(key, 0.0) + float(val.item())