Spaces:
Runtime error
Runtime error
File size: 5,757 Bytes
a5f8a35 |
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 164 165 166 167 168 169 170 171 172 |
import argparse
import json
import os
import random
from typing import Any, Dict, List
from loguru import logger
import torch
from torch.utils.data import DataLoader, DistributedSampler
from torch.nn.utils.rnn import pad_sequence
from tqdm import tqdm
import wordsegment as ws
from virtex.config import Config
from virtex.data import ZeroShotDataset
from virtex.data.tokenizers import SentencePieceBPETokenizer
from virtex.factories import TokenizerFactory, VisualBackboneFactory,TextualHeadFactory
from virtex.utils.checkpointing import CheckpointManager
from virtex.utils.common import common_parser
from virtex.utils.metrics import TopkAccuracy
import virtex.utils.distributed as dist
#importing classifier
from virtex.models.zero_shot_classification_eval import ZeroShotClassifier
ws.load()
# fmt: off
parser = common_parser(
description="""Run image captioning inference on a pretrained model, and/or
evaluate pretrained model on COCO Captions val2017 split."""
)
parser.add_argument(
"--data-root", default=None,
help="""Path to a directory containing image files to generate captions for imagenet.
Default: COCO val2017 image directory as expected relative to project root."""
)
parser.add_argument(
"--checkpoint-path", required=False,
help="Path to load checkpoint and run captioning evaluation."
)
parser.add_argument(
"--output", default=None,
help="Path to save predictions as a JSON file."
)
parser.add_argument(
"--calc-metrics", action="store_true",
help="""Calculate CIDEr and SPICE metrics using ground truth COCO Captions.
This flag should not be set when running inference on arbitrary images."""
)
parser.add_argument(
"--idx_label_dict", default=None, required=False,
help="""a dictionary that maps from lable index to label string for classification"""
)
parser.add_argument(
"--is_redcaps", default=None, required=False,
help="""a dictionary that maps from lable index to label string for"""
)
parser.add_argument(
"--prompt_cls_sos", default=None, required=False,
help="""a dictionary that maps from lable index to label string for"""
)
parser.add_argument(
"--prompt_sos_eos", default=None, required=False,
help="""a dictionary that maps from lable index to label string for"""
)
# fmt: on
print("###########")
print(os.getcwd() )
print("###########")
tokenizer = SentencePieceBPETokenizer("datasets_1/vocab/common_32k.model")
def main(_A: argparse.Namespace):
if _A.num_gpus_per_machine == 0:
# Set device as CPU if num_gpus_per_machine = 0.
device = torch.device("cpu")
else:
# Get the current device (this will be zero here by default).
device = torch.cuda.current_device()
_C = Config(_A.config, _A.config_override)
#tokenizer = TokenizerFactory.from_config(_C)
if _A.data_root is None:
_A.data_root = os.path.join(_C.DATA.ROOT, "val2017")
if _A.is_redcaps == 1:
model_dataset = 'redcaps'
else:
model_dataset = 'gcc or sbu'
print(_A.idx_label_dict)
val_dataset = ZeroShotDataset(data_root=_A.data_root,
split="test/",
label_map=_A.idx_label_dict,
tokenizer=tokenizer,
prompt_cls_sos=_A.prompt_cls_sos.replace("_", " "),
prompt_sos_eos=_A.prompt_sos_eos.replace("_", " "))
val_dataloader = DataLoader(
val_dataset,
batch_size= _C.OPTIM.BATCH_SIZE // dist.get_world_size(),
num_workers=_A.cpu_workers,
sampler=DistributedSampler(
val_dataset,
num_replicas=dist.get_world_size(),
rank=dist.get_rank(),
),
pin_memory=True,
drop_last=False,
collate_fn=val_dataset.collate_fn,
)
# Initialize model from a checkpoint
visual = VisualBackboneFactory.from_config(_C)
textual = TextualHeadFactory.from_config(_C)
model = ZeroShotClassifier(visual,textual)
ITERATION = CheckpointManager(model=model).load(_A.checkpoint_path)
model.to(device).eval()
## setup distributed training
if dist.get_world_size() > 1:
dist.synchronize()
model = nn.parallel.DistributedDataParallel(
model, device_ids=[device], find_unused_parameters=True
)
top_1 = TopkAccuracy(top_k=1)
top_5 = TopkAccuracy(top_k=5)
batch_num = 0
for val_iteration, val_batch in tqdm(enumerate(val_dataloader, start=1)):
val_batch["image"] = val_batch["image"].to(device)
val_batch["caption_tokens"] = val_batch["caption_tokens"].to(device)
val_batch["noitpac_tokens"] = val_batch["noitpac_tokens"] .to(device)
val_batch["caption_lengths"] = val_batch["caption_lengths"].to(device)
val_batch["label"] = val_batch["label"].to(device)
with torch.no_grad():
classification_losses = model(val_batch)
batch_num+=1
top_1(classification_losses, val_batch["label"])
top_1_acc = top_1.get_metric(reset=False)
dist.average_across_processes(top_1_acc)
top_5(classification_losses, val_batch["label"])
top_5_acc = top_5.get_metric(reset=False)
dist.average_across_processes(top_5_acc)
logger.info(f"Iter: {val_iteration} | Top-1 accuracy: {top_1_acc} | Top-5 accuracy: {top_5_acc}")
if __name__ == "__main__":
_A = parser.parse_args()
#if _A.num_gpus_per_machine > 1:
# raise ValueError("Using multiple GPUs is not supported for this script.")
# No distributed training here, just a single process.
main(_A)
|