File size: 6,131 Bytes
002bd9b |
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 |
# TODO: extract images from refcoco series
import sys
sys.path.append(".")
import base64
import io
import json
import logging
import os
import os.path as osp
import datasets
import hydra
import numpy as np
import tqdm
from hydra.core.hydra_config import HydraConfig
from hydra.core.utils import configure_log
from omegaconf import DictConfig, OmegaConf
from PIL import Image
import pycocotools.mask
from utils.git_utils import TSVWriter
from src.arguments import Arguments, global_setup
import logging
from hydra.utils import instantiate
from transformers import set_seed, AutoTokenizer
from datasets import interleave_datasets, concatenate_datasets
import torch
from src.train import prepare_datasets
logger = logging.getLogger(__name__)
@hydra.main(version_base="1.3", config_path="../../src/conf", config_name="conf")
def main(args: Arguments):
logger.warning(f"Turn no_cuda = True.")
args.training.no_cuda = True
# NOTE: ddp is initialized in _setup_devices class in `transformers/training_args.py`
args, training_args, _ = global_setup(args)
dummy_data = torch.tensor(training_args.process_index)
dummy_data_ls = [torch.tensor(1) for _ in range(training_args.world_size)]
torch.distributed.all_gather(dummy_data_ls, dummy_data)
logger.info(f"Try gloo ddp: rank {training_args.process_index} dummy_data_ls: {dummy_data_ls}")
# Set seed before initializing model.
set_seed(args.training.seed)
# Initialize our dataset and prepare it
train_dataset, eval_dataset = prepare_datasets(args)
if len(eval_dataset) > 1:
raise ValueError(f"Only support one eval dataset, but got {len(eval_dataset)}. args: {args.eval_data}")
logger.info(f"world_size {training_args.world_size}")
logger.info(f"process_index {training_args.process_index}")
logger.info(f"local_process_index {training_args.local_process_index}")
if training_args.world_size > 1:
from datasets.distributed import split_dataset_by_node
# NOTE: to debug ddp code
# if training_args.should_log:
# breakpoint()
# torch.distributed.barrier()
logger.info(f"Splitting dataset by node {training_args.process_index}/{training_args.world_size}")
if train_dataset is not None:
train_dataset = split_dataset_by_node(train_dataset, training_args.process_index, training_args.world_size)
if eval_dataset is not None:
for eval_dataset_k, eval_dataset_v in eval_dataset.items():
eval_dataset[eval_dataset_k] = split_dataset_by_node(
eval_dataset_v, training_args.process_index, training_args.world_size
)
logger.info(f"[{training_args.process_index}/{training_args.world_size}]: Train Dataset: {train_dataset}")
logger.info(f"[{training_args.process_index}/{training_args.world_size}]: Eval Dataset: {eval_dataset}")
# NOTE: According to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb,
# we use alternatively GPT2TokenizerFast.
tokenizer = AutoTokenizer.from_pretrained("gpt2")
process_dataset(train_dataset, "train", training_args, tokenizer)
for eval_dataset_k, eval_dataset_v in eval_dataset.items():
process_dataset(eval_dataset_v, f"eval-{eval_dataset_k}", training_args, tokenizer)
def process_dataset(dataset, split_name, training_args, tokenizer):
if dataset is None:
logger.warning(
f"[{training_args.process_index}/{training_args.world_size}]: {split_name} is None, skip processing"
)
return
region_cnt = 0
sample_cnt = 0
sent_cnt = 0
token_cnt = 0
word_cnt = 0
tbar = tqdm.trange(len(dataset), desc=f"Processing {split_name}")
for sample_idx in tbar:
try:
sample = dataset[sample_idx]
except Exception as e:
logger.error(f"Error when processing {split_name} sample_idx {sample_idx}: {e}")
continue
for annot_idx, region in enumerate(sample["regions"]):
cap = region["phrases"]
for cap_ in cap:
# NOTE: https://huggingface.co/docs/transformers/main_classes/tokenizer#transformers.PreTrainedTokenizer.tokenize
token_cnt += len(tokenizer.tokenize(cap_))
# NOTE: https://huggingface.co/learn/nlp-course/chapter6/4#pre-tokenization
word_cnt += len(tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(cap_))
sent_cnt += len(cap)
region_cnt += 1
sample_cnt += 1
tbar.set_description(
f"[{training_args.process_index}/{training_args.world_size}]: Already processing {sample_cnt} samples, {region_cnt} regions, {sent_cnt} sentences, and {token_cnt} tokens."
)
logger.info(
f"[{training_args.process_index}/{training_args.world_size}]: Total samples: {sample_cnt}, total regions: {region_cnt}, total sents: {sent_cnt}, total tokens: {token_cnt}"
)
def _gather_cnts(cnt, training_args, name=""):
if training_args.world_size > 1:
cnt_ls = [torch.tensor(0) for _ in range(training_args.world_size)]
torch.distributed.all_gather(cnt_ls, torch.tensor(cnt))
cnt = sum(cnt_ls)
if training_args.process_index == 0:
logger.info(f"[{training_args.process_index}/{training_args.world_size}]: {name}_cnt_ls: {cnt}")
return cnt
all_sample_cnt = _gather_cnts(sample_cnt, training_args, "sample")
all_region_cnt = _gather_cnts(region_cnt, training_args, "region")
all_sent_cnt = _gather_cnts(sent_cnt, training_args, "sent")
all_token_cnt = _gather_cnts(token_cnt, training_args, "token")
all_word_cnt = _gather_cnts(word_cnt, training_args, "word")
if training_args.process_index == 0:
logger.info(
f"[FULL]: split name: {split_name}, total samples: {all_sample_cnt}, total regions: {all_region_cnt}, total sents: {all_sent_cnt}, total tokens: {all_token_cnt}, total words: {all_word_cnt}"
)
if __name__ == "__main__":
main()
|