Vision-CAIR
commited on
Commit
•
1119ef9
1
Parent(s):
ad32631
Upload folder using huggingface_hub
Browse files
blip2.py
CHANGED
@@ -15,8 +15,7 @@ import torch.nn as nn
|
|
15 |
import torch.distributed as dist
|
16 |
import torch.nn.functional as F
|
17 |
|
18 |
-
|
19 |
-
from .dist_utils import download_cached_file
|
20 |
from .utils import is_url
|
21 |
from .logger import MetricLogger
|
22 |
from .base_model import BaseModel
|
@@ -166,8 +165,8 @@ def compute_sim_matrix(model, data_loader, **kwargs):
|
|
166 |
(len(data_loader.dataset.image), len(texts)), -100.0
|
167 |
).to(model.device)
|
168 |
|
169 |
-
num_tasks =
|
170 |
-
rank =
|
171 |
step = sims_matrix.size(0) // num_tasks + 1
|
172 |
start = rank * step
|
173 |
end = min(sims_matrix.size(0), start + step)
|
@@ -205,7 +204,7 @@ def compute_sim_matrix(model, data_loader, **kwargs):
|
|
205 |
).float()
|
206 |
score_matrix_t2i[start + i, topk_idx] = score + topk_sim
|
207 |
|
208 |
-
if
|
209 |
dist.barrier()
|
210 |
torch.distributed.all_reduce(
|
211 |
score_matrix_i2t, op=torch.distributed.ReduceOp.SUM
|
|
|
15 |
import torch.distributed as dist
|
16 |
import torch.nn.functional as F
|
17 |
|
18 |
+
from .dist_utils import download_cached_file,get_world_size,get_rank,is_dist_avail_and_initialized
|
|
|
19 |
from .utils import is_url
|
20 |
from .logger import MetricLogger
|
21 |
from .base_model import BaseModel
|
|
|
165 |
(len(data_loader.dataset.image), len(texts)), -100.0
|
166 |
).to(model.device)
|
167 |
|
168 |
+
num_tasks = get_world_size()
|
169 |
+
rank = get_rank()
|
170 |
step = sims_matrix.size(0) // num_tasks + 1
|
171 |
start = rank * step
|
172 |
end = min(sims_matrix.size(0), start + step)
|
|
|
204 |
).float()
|
205 |
score_matrix_t2i[start + i, topk_idx] = score + topk_sim
|
206 |
|
207 |
+
if is_dist_avail_and_initialized():
|
208 |
dist.barrier()
|
209 |
torch.distributed.all_reduce(
|
210 |
score_matrix_i2t, op=torch.distributed.ReduceOp.SUM
|