Spaces:
Runtime error
Runtime error
File size: 6,617 Bytes
8fc2b4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
# Copyright (c) Aishwarya Kamath & Nicolas Carion. Licensed under the Apache License 2.0. All Rights Reserved
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Misc functions, including distributed helpers.
Mostly copy-paste from torchvision references.
"""
import os
import subprocess
from typing import Any, Dict, List, Optional
import torch
import torchvision
from torch import Tensor
def get_sha():
cwd = os.path.dirname(os.path.abspath(__file__))
def _run(command):
return subprocess.check_output(command, cwd=cwd).decode("ascii").strip()
sha = "N/A"
diff = "clean"
branch = "N/A"
try:
sha = _run(["git", "rev-parse", "HEAD"])
subprocess.check_output(["git", "diff"], cwd=cwd)
diff = _run(["git", "diff-index", "HEAD"])
diff = "has uncommited changes" if diff else "clean"
branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"])
except Exception:
pass
message = f"sha: {sha}, status: {diff}, branch: {branch}"
return message
def collate_fn(do_round, batch):
batch = list(zip(*batch))
final_batch = {}
final_batch["samples"] = NestedTensor.from_tensor_list(batch[0], do_round)
final_batch["targets"] = batch[1]
if "positive_map" in batch[1][0]:
# we batch the positive maps here
# Since in general each batch element will have a different number of boxes,
# we collapse a single batch dimension to avoid padding. This is sufficient for our purposes.
max_len = max([v["positive_map"].shape[1] for v in batch[1]])
nb_boxes = sum([v["positive_map"].shape[0] for v in batch[1]])
batched_pos_map = torch.zeros((nb_boxes, max_len), dtype=torch.bool)
cur_count = 0
for v in batch[1]:
cur_pos = v["positive_map"]
batched_pos_map[cur_count : cur_count + len(cur_pos), : cur_pos.shape[1]] = cur_pos
cur_count += len(cur_pos)
assert cur_count == len(batched_pos_map)
# assert batched_pos_map.sum().item() == sum([v["positive_map"].sum().item() for v in batch[1]])
final_batch["positive_map"] = batched_pos_map.float()
if "positive_map_eval" in batch[1][0]:
# we batch the positive maps here
# Since in general each batch element will have a different number of boxes,
# we collapse a single batch dimension to avoid padding. This is sufficient for our purposes.
max_len = max([v["positive_map_eval"].shape[1] for v in batch[1]])
nb_boxes = sum([v["positive_map_eval"].shape[0] for v in batch[1]])
batched_pos_map = torch.zeros((nb_boxes, max_len), dtype=torch.bool)
cur_count = 0
for v in batch[1]:
cur_pos = v["positive_map_eval"]
batched_pos_map[cur_count : cur_count + len(cur_pos), : cur_pos.shape[1]] = cur_pos
cur_count += len(cur_pos)
assert cur_count == len(batched_pos_map)
# assert batched_pos_map.sum().item() == sum([v["positive_map"].sum().item() for v in batch[1]])
final_batch["positive_map_eval"] = batched_pos_map.float()
if "answer" in batch[1][0] or "answer_type" in batch[1][0]:
answers = {}
for f in batch[1][0].keys():
if "answer" not in f:
continue
answers[f] = torch.stack([b[f] for b in batch[1]])
final_batch["answers"] = answers
return final_batch
class NestedTensor(object):
def __init__(self, tensors, mask):
self.tensors = tensors
self.mask = mask
def to(self, *args, **kwargs):
cast_tensor = self.tensors.to(*args, **kwargs)
cast_mask = self.mask.to(*args, **kwargs) if self.mask is not None else None
return type(self)(cast_tensor, cast_mask)
def decompose(self):
return self.tensors, self.mask
@classmethod
def from_tensor_list(cls, tensor_list, do_round=False):
# TODO make this more general
if tensor_list[0].ndim == 3:
# TODO make it support different-sized images
max_size = tuple(max(s) for s in zip(*[img.shape for img in tensor_list]))
# min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
batch_shape = (len(tensor_list),) + max_size
b, c, h, w = batch_shape
if do_round:
# Round to an even size to avoid rounding issues in fpn
p = 128
h = h if h % p == 0 else (h // p + 1) * p
w = w if w % p == 0 else (w // p + 1) * p
batch_shape = b, c, h, w
dtype = tensor_list[0].dtype
device = tensor_list[0].device
tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
for img, pad_img, m in zip(tensor_list, tensor, mask):
pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
m[: img.shape[1], : img.shape[2]] = False
else:
raise ValueError("not supported")
return cls(tensor, mask)
def __repr__(self):
return repr(self.tensors)
def interpolate(input, size=None, scale_factor=None, mode="nearest", align_corners=None):
# type: (Tensor, Optional[List[int]], Optional[float], str, Optional[bool]) -> Tensor
"""
Equivalent to nn.functional.interpolate, but with support for empty channel sizes.
"""
if input.numel() > 0:
return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)
assert input.shape[0] != 0 or input.shape[1] != 0, "At least one of the two first dimensions must be non zero"
if input.shape[1] == 0:
# Pytorch doesn't support null dimension on the channel dimension, so we transpose to fake a null batch dim
return torch.nn.functional.interpolate(input.transpose(0, 1), size, scale_factor, mode, align_corners).transpose(0, 1)
# empty batch dimension is now supported in pytorch
return torch.nn.functional.interpolate(input, size, scale_factor, mode, align_corners)
def targets_to(targets: List[Dict[str, Any]], device):
"""Moves the target dicts to the given device."""
excluded_keys = [
"questionId",
"tokens_positive",
"tokens",
"dataset_name",
"sentence_id",
"original_img_id",
"nb_eval",
"task_id",
"original_id",
]
return [{k: v.to(device) if k not in excluded_keys else v for k, v in t.items() if k != "caption"} for t in targets]
|