Delete utils_.py
Browse files
utils_.py
DELETED
|
@@ -1,161 +0,0 @@
|
|
| 1 |
-
import os
|
| 2 |
-
import math
|
| 3 |
-
import numpy as np
|
| 4 |
-
import torch
|
| 5 |
-
import torch.distributed as dist
|
| 6 |
-
import torchvision.transforms as T
|
| 7 |
-
from decord import VideoReader, cpu
|
| 8 |
-
from PIL import Image
|
| 9 |
-
from transformers import AutoConfig
|
| 10 |
-
from torchvision.transforms.functional import InterpolationMode
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
| 14 |
-
IMAGENET_STD = (0.229, 0.224, 0.225)
|
| 15 |
-
|
| 16 |
-
def load_image(image_file, input_size=448, max_num=12, upscale=False):
|
| 17 |
-
image = Image.open(image_file).convert('RGB')
|
| 18 |
-
if upscale:
|
| 19 |
-
image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
|
| 20 |
-
transform = build_transform(input_size=input_size)
|
| 21 |
-
images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
|
| 22 |
-
pixel_values = [transform(image) for image in images]
|
| 23 |
-
pixel_values = torch.stack(pixel_values)
|
| 24 |
-
return pixel_values
|
| 25 |
-
|
| 26 |
-
def build_transform(input_size):
|
| 27 |
-
MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
|
| 28 |
-
transform = T.Compose([
|
| 29 |
-
T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
|
| 30 |
-
T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
|
| 31 |
-
T.ToTensor(),
|
| 32 |
-
T.Normalize(mean=MEAN, std=STD)
|
| 33 |
-
])
|
| 34 |
-
return transform
|
| 35 |
-
|
| 36 |
-
def get_rank_and_world_size():
|
| 37 |
-
rank = int(os.environ.get('RANK', 0))
|
| 38 |
-
world_size = int(os.environ.get('WORLD_SIZE', 1))
|
| 39 |
-
return rank, world_size
|
| 40 |
-
|
| 41 |
-
def get_local_rank_and_local_world_size():
|
| 42 |
-
if not dist.is_available():
|
| 43 |
-
return 0, 1
|
| 44 |
-
if not dist.is_initialized():
|
| 45 |
-
return 0, 1
|
| 46 |
-
|
| 47 |
-
if 'SLURM_LOCALID' in os.environ:
|
| 48 |
-
local_rank = int(os.environ['SLURM_LOCALID'])
|
| 49 |
-
local_world_size = int(os.environ['SLURM_NTASKS_PER_NODE'])
|
| 50 |
-
return local_rank, local_world_size
|
| 51 |
-
|
| 52 |
-
if 'LOCAL_RANK' in os.environ and 'LOCAL_WORLD_SIZE' in os.environ:
|
| 53 |
-
return int(os.environ['LOCAL_RANK']), int(os.environ['LOCAL_WORLD_SIZE'])
|
| 54 |
-
|
| 55 |
-
raise NotImplementedError(
|
| 56 |
-
"Fail to get local_rank and local_world_size! "
|
| 57 |
-
"Please ensure that you set the environment variable "
|
| 58 |
-
"`LOCAL_RANK` and `LOCAL_WORLD_SIZE`"
|
| 59 |
-
)
|
| 60 |
-
|
| 61 |
-
def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
|
| 62 |
-
best_ratio_diff = float('inf')
|
| 63 |
-
best_ratio = (1, 1)
|
| 64 |
-
area = width * height
|
| 65 |
-
for ratio in target_ratios:
|
| 66 |
-
target_aspect_ratio = ratio[0] / ratio[1]
|
| 67 |
-
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
|
| 68 |
-
if ratio_diff < best_ratio_diff:
|
| 69 |
-
best_ratio_diff = ratio_diff
|
| 70 |
-
best_ratio = ratio
|
| 71 |
-
elif ratio_diff == best_ratio_diff:
|
| 72 |
-
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
|
| 73 |
-
best_ratio = ratio
|
| 74 |
-
return best_ratio
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
|
| 78 |
-
orig_width, orig_height = image.size
|
| 79 |
-
aspect_ratio = orig_width / orig_height
|
| 80 |
-
|
| 81 |
-
# calculate the existing image aspect ratio
|
| 82 |
-
target_ratios = set(
|
| 83 |
-
(i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
|
| 84 |
-
i * j <= max_num and i * j >= min_num)
|
| 85 |
-
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
|
| 86 |
-
|
| 87 |
-
# find the closest aspect ratio to the target
|
| 88 |
-
target_aspect_ratio = find_closest_aspect_ratio(
|
| 89 |
-
aspect_ratio, target_ratios, orig_width, orig_height, image_size)
|
| 90 |
-
|
| 91 |
-
# calculate the target width and height
|
| 92 |
-
target_width = image_size * target_aspect_ratio[0]
|
| 93 |
-
target_height = image_size * target_aspect_ratio[1]
|
| 94 |
-
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
|
| 95 |
-
|
| 96 |
-
# resize the image
|
| 97 |
-
resized_img = image.resize((target_width, target_height))
|
| 98 |
-
processed_images = []
|
| 99 |
-
for i in range(blocks):
|
| 100 |
-
box = (
|
| 101 |
-
(i % (target_width // image_size)) * image_size,
|
| 102 |
-
(i // (target_width // image_size)) * image_size,
|
| 103 |
-
((i % (target_width // image_size)) + 1) * image_size,
|
| 104 |
-
((i // (target_width // image_size)) + 1) * image_size
|
| 105 |
-
)
|
| 106 |
-
# split the image
|
| 107 |
-
split_img = resized_img.crop(box)
|
| 108 |
-
processed_images.append(split_img)
|
| 109 |
-
assert len(processed_images) == blocks
|
| 110 |
-
if use_thumbnail and len(processed_images) != 1:
|
| 111 |
-
thumbnail_img = image.resize((image_size, image_size))
|
| 112 |
-
processed_images.append(thumbnail_img)
|
| 113 |
-
return processed_images
|
| 114 |
-
|
| 115 |
-
def split_model(model_path):
|
| 116 |
-
num_gpus_per_node = torch.cuda.device_count()
|
| 117 |
-
rank, world_size = get_rank_and_world_size()
|
| 118 |
-
try:
|
| 119 |
-
local_rank, local_world_size = get_local_rank_and_local_world_size()
|
| 120 |
-
except:
|
| 121 |
-
local_rank = rank
|
| 122 |
-
|
| 123 |
-
if 'GPUS_PER_PROCESS' in os.environ:
|
| 124 |
-
gpus_per_process = int(os.environ['GPUS_PER_PROCESS'])
|
| 125 |
-
else:
|
| 126 |
-
gpus_per_process = 8 # default to use 8 GPUs for one model
|
| 127 |
-
gpus_per_process = min(gpus_per_process, num_gpus_per_node // local_world_size)
|
| 128 |
-
start_gpu = local_rank * gpus_per_process
|
| 129 |
-
end_gpu = start_gpu + gpus_per_process
|
| 130 |
-
|
| 131 |
-
assert end_gpu <= num_gpus_per_node, f"Process {local_rank} tries to access GPU {end_gpu}, " \
|
| 132 |
-
f"but only {num_gpus_per_node} GPUs are available per node."
|
| 133 |
-
|
| 134 |
-
visible_devices = list(range(start_gpu, end_gpu))
|
| 135 |
-
|
| 136 |
-
device_map = {}
|
| 137 |
-
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
|
| 138 |
-
|
| 139 |
-
num_gpus_for_vit = 0.5
|
| 140 |
-
num_layers = config.llm_config.num_hidden_layers
|
| 141 |
-
num_layers_per_gpu = math.ceil(num_layers / (len(visible_devices) - num_gpus_for_vit))
|
| 142 |
-
num_layers_per_gpu = [num_layers_per_gpu] * len(visible_devices)
|
| 143 |
-
num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
|
| 144 |
-
|
| 145 |
-
layer_cnt = 0
|
| 146 |
-
for i, num_layer in enumerate(num_layers_per_gpu):
|
| 147 |
-
for j in range(num_layer):
|
| 148 |
-
device_map[f'language_model.model.layers.{layer_cnt}'] = visible_devices[i]
|
| 149 |
-
layer_cnt += 1
|
| 150 |
-
device_map['vision_model'] = visible_devices[0]
|
| 151 |
-
device_map['mlp1'] = visible_devices[0]
|
| 152 |
-
device_map['language_model.model.tok_embeddings'] = visible_devices[0]
|
| 153 |
-
device_map['language_model.model.embed_tokens'] = visible_devices[0]
|
| 154 |
-
device_map['language_model.output'] = visible_devices[0]
|
| 155 |
-
device_map['language_model.model.norm'] = visible_devices[0]
|
| 156 |
-
device_map['language_model.model.rotary_emb'] = visible_devices[0]
|
| 157 |
-
device_map['language_model.lm_head'] = visible_devices[0]
|
| 158 |
-
device_map[f'language_model.model.layers.{num_layers - 1}'] = visible_devices[0]
|
| 159 |
-
|
| 160 |
-
return device_map, visible_devices
|
| 161 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|