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 return lr_lambda
@torch.no_grad()
def _embed_drone_queries( def _embed_drone_queries(
model: AsymmetricEncoder, model: AsymmetricEncoder,
train_ds: GTAUAVDataset, train_ds: GTAUAVDataset,
@@ -153,24 +153,19 @@ def _embed_drone_queries(
collate_fn=collate_drone_query, collate_fn=collate_drone_query,
pin_memory=True, pin_memory=True,
) )
all_embs: list[torch.Tensor] = []
with torch.inference_mode(): embs: list[torch.Tensor] = []
for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False):
drone_img = batch["drone_img"].to(device, non_blocking=True) drone_img = batch["drone_img"].to(device, non_blocking=True)
altitude = batch.get("altitude") q = model.encode_query(
if altitude is not None: drone_img,
altitude = altitude.to(device, non_blocking=True) batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
kwargs: dict[str, Any] = {"drone_img": drone_img, "altitude": altitude} )
if not getattr(model, "baseline_mode", False): embs.append(q.cpu())
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())
if was_training: if was_training:
model.train() model.train()
return torch.cat(all_embs, dim=0) return torch.cat(embs, dim=0)
class Trainer: class Trainer:
@@ -743,13 +738,10 @@ class Trainer:
drone_img = batch["drone_img"].to(device, non_blocking=True) drone_img = batch["drone_img"].to(device, non_blocking=True)
sat_img = batch["sat_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): with autocast(device_type="cuda", enabled=self.hardware_cfg.use_amp):
if baseline_mode: 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: else:
embeddings = self.model( embeddings = self.model(
drone_img=drone_img, drone_img=drone_img,
@@ -759,8 +751,7 @@ class Trainer:
caption_l3=batch["caption_l3"], caption_l3=batch["caption_l3"],
sat_caption_l1=batch["sat_caption_l1"], sat_caption_l1=batch["sat_caption_l1"],
sat_caption_l2=batch["sat_caption_l2"], sat_caption_l2=batch["sat_caption_l2"],
sat_caption_l3=batch["sat_caption_l3"], sat_caption_l3=batch["sat_caption_l3"]
altitude=altitude,
) )
queue_neg = self.neg_bank.get_queue() if self.neg_bank is not None else None queue_neg = self.neg_bank.get_queue() if self.neg_bank is not None else None