| 115 | ) |
| 116 | |
| 117 | def get_cap_generator(): |
| 118 | while True: |
| 119 | for data in tqdm( |
| 120 | loader, |
| 121 | disable=not (rank == 0), |
| 122 | initial=0, |
| 123 | desc=f"generate_captions, for iters {iterations}", |
| 124 | ): |
| 125 | if has_text(args): |
| 126 | _cap_feats, _cap = ( |
| 127 | data["caption_feature"].to(device), |
| 128 | data["caption"], |
| 129 | ) |
| 130 | B, N, T, C = _cap_feats.shape # each image has N captions |
| 131 | _p = random.randint(0, N - 1) |
| 132 | yield _cap_feats[:, _p], [_cap[i][_p] for i in range(len(_cap))] |
| 133 | else: |
| 134 | raise NotImplementedError("current dataset doesnt have captions") |
| 135 | |
| 136 | if has_text(args): |
| 137 | cap_dg = get_cap_generator() |