Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

claude_refactor_v3: fix extra lines in trainer_new

Browse Source
This commit is contained in:
pikaliov
2026-05-05 12:37:20 +03:00
parent 248bd331d2
commit 4cb5ab97d8
1 changed files with 14 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
src/training/trainer_new.py
View File

@@ -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