| | |
| | |
| | |
| | |
| | |
| |
|
| | 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) |
| |
|
