From 4cb5ab97d8ffa9dcef6db285add3923193fb0d12 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 5 May 2026 12:37:20 +0300 Subject: [PATCH] claude_refactor_v3: fix extra lines in trainer_new --- src/training/trainer_new.py | 37 ++++++++++++++----------------------- 1 file changed, 14 insertions(+), 23 deletions(-) diff --git a/src/training/trainer_new.py b/src/training/trainer_new.py index 268a281..3086d53 100644 --- a/src/training/trainer_new.py +++ b/src/training/trainer_new.py @@ -128,7 +128,7 @@ def _cosine_warmup_schedule(warmup_steps: int, total_steps: int): return lr_lambda - +@torch.no_grad() def _embed_drone_queries( model: AsymmetricEncoder, train_ds: GTAUAVDataset, @@ -153,24 +153,19 @@ def _embed_drone_queries( collate_fn=collate_drone_query, pin_memory=True, ) - all_embs: list[torch.Tensor] = [] - with torch.inference_mode(): - for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): - drone_img = batch["drone_img"].to(device, non_blocking=True) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) - kwargs: dict[str, Any] = {"drone_img": drone_img, "altitude": altitude} - if not getattr(model, "baseline_mode", False): - kwargs["caption_l1"] = batch["caption_l1"] - kwargs["caption_l2"] = batch["caption_l2"] - kwargs["caption_l3"] = batch["caption_l3"] - with autocast(device_type="cuda", enabled=True): - emb = model.encode_query(**kwargs) - all_embs.append(emb.cpu()) + + embs: list[torch.Tensor] = [] + for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): + drone_img = batch["drone_img"].to(device, non_blocking=True) + q = model.encode_query( + drone_img, + batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + ) + embs.append(q.cpu()) + if was_training: model.train() - return torch.cat(all_embs, dim=0) + return torch.cat(embs, dim=0) class Trainer: @@ -743,13 +738,10 @@ class Trainer: drone_img = batch["drone_img"].to(device, non_blocking=True) sat_img = batch["sat_img"].to(device, non_blocking=True) - altitude = batch.get("altitude") - if altitude is not None: - altitude = altitude.to(device, non_blocking=True) with autocast(device_type="cuda", enabled=self.hardware_cfg.use_amp): if baseline_mode: - embeddings = self.model(drone_img=drone_img, sat_img=sat_img, altitude=altitude) + embeddings = self.model(drone_img=drone_img, sat_img=sat_img) else: embeddings = self.model( drone_img=drone_img, @@ -759,8 +751,7 @@ class Trainer: caption_l3=batch["caption_l3"], sat_caption_l1=batch["sat_caption_l1"], sat_caption_l2=batch["sat_caption_l2"], - sat_caption_l3=batch["sat_caption_l3"], - altitude=altitude, + sat_caption_l3=batch["sat_caption_l3"] ) queue_neg = self.neg_bank.get_queue() if self.neg_bank is not None else None