File size: 4,053 Bytes
866ee56 |
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 |
# --------------------------------------------------------
# InternVL
# Copyright (c) 2023 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
import argparse
import time
import torch
from mmcv.cnn import get_model_complexity_info
from mmcv.cnn.utils.flops_counter import flops_to_string, params_to_string
from models.intern_vit_6b import InternViT6B
from tqdm import tqdm
parser = argparse.ArgumentParser(description='Hyperparams')
parser.add_argument('config', nargs='?', type=str, default=None)
args = parser.parse_args()
configs = {
'a': {
'embed_dim': 3968,
'num_heads': 62,
'mlp_ratio': 4,
'depth': 32
},
'e': {
'embed_dim': 3200,
'num_heads': 50,
'mlp_ratio': 4,
'depth': 48
},
'f': {
'embed_dim': 3200,
'num_heads': 25,
'mlp_ratio': 4,
'depth': 48
},
'g': {
'embed_dim': 2496,
'num_heads': 39,
'mlp_ratio': 8,
'depth': 48
},
'i': {
'embed_dim': 2816,
'num_heads': 44,
'mlp_ratio': 4,
'depth': 64
},
'm': {
'embed_dim': 2496,
'num_heads': 39,
'mlp_ratio': 4,
'depth': 80
},
}
def sa_flops(h, w, dim):
return 2 * h * w * h * w * dim
def get_flops(model, input_shape):
flops, params = get_model_complexity_info(model,
input_shape,
as_strings=False)
_, H, W = input_shape
print(flops, params)
for i in range(model.depth):
flops += sa_flops(H // model.patch_size, W // model.patch_size,
model.embed_dim)
return flops_to_string(flops), params_to_string(params)
if __name__ == '__main__':
input_shape = (3, 224, 224)
config = configs[args.config]
print(config)
model = InternViT6B(in_chans=3,
patch_size=14,
img_size=224,
pretrain_size=224,
qkv_bias=False,
drop_path_rate=0.0,
embed_dim=config['embed_dim'],
num_heads=config['num_heads'],
mlp_ratio=config['mlp_ratio'],
init_values=0.1,
qk_normalization=True,
depth=config['depth'],
use_flash_attn=True,
with_cp=True,
freeze_vit=True,
cls_target='cls_patch_concat',
num_classes=0,
attn_pool_num_heads=16,
clip_embed_dim=768,
head_norm_type='bn').to(torch.bfloat16)
for k, v in model.named_parameters():
v.requires_grad = True
if torch.cuda.is_available():
model.cuda()
model.eval()
flops, params = get_flops(model, input_shape)
split_line = '=' * 30
print(f'{split_line}\nInput shape: {input_shape}\n'
f'Flops: {flops}\nParams: {params}\n{split_line}')
print('!!!Please be cautious if you use the results in papers. '
'You may need to check if all ops are supported and verify that the '
'flops computation is correct.')
image = torch.rand(128, 3, 224, 224).to(torch.bfloat16).cuda()
torch.cuda.synchronize()
start_time = time.time()
with torch.no_grad():
for i in tqdm(range(10)):
out = model(image)
torch.cuda.synchronize()
end_time = time.time()
print('warmup time: ', end_time - start_time)
torch.cuda.synchronize()
start_time = time.time()
with torch.no_grad():
for i in tqdm(range(50)):
out = model(image)
torch.cuda.synchronize()
end_time = time.time()
print('using time: ', (end_time - start_time))
print('FPS: ', 50 * 128 / (end_time - start_time))
print(config)
|