File size: 3,665 Bytes
8b3f571 | 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 | _base_=['./default_runtime.py']
nb_code = 2048
unit_length=4
audio_unit_length=8
find_unused_parameters=True
model_size = '3B'
mean_path='data/mean_272.npy'
std_path='data/std_272.npy'
traj_mean_path='data/traj_mean.npy'
traj_std_path='data/traj_std.npy'
vqvae_cfg = dict(
type='VQVAE',
init_cfg=dict(
type='Pretrained',
checkpoint='work_dirs/tokenizer_r_fsq_2048cd_4res_200e/epoch_200.pth'),
decoder_cfg=dict(
activation='relu',
depth=3, dilation_growth_rate=3,
down_t=2, input_emb_width=272,
norm=None, output_emb_width=512,
type='DecoderAttn', width=512),
encoder_cfg=dict(
activation='relu',
depth=3, dilation_growth_rate=3,
down_t=2, input_emb_width=272,
norm=None, output_emb_width=512,
stride_t=2, type='EncoderAttn', width=512),
quantizer_cfg=dict(
dim=512,
levels=[8, 8, 6, 5],
num_quantizers=4,
quantize_dropout_prob=0.2,
type='ResidualFSQ'),
)
llama_configs = {
"70M": dict(n_layer=8, n_head=8, n_embd=512),
"111M": dict(n_layer=12, n_head=12, n_embd=768),
"343M": dict(n_layer=24, n_head=16, n_embd=1024),
"775M": dict(n_layer=36, n_head=20, n_embd=1280),
"1B": dict(n_layer=48, n_head=24, n_embd=1536),
"3B": dict(n_layer=24, n_head=32, n_embd=3200),
"7B": dict(n_layer=32, n_head=32, n_embd=4096),
}
# model
model=dict(
type='AnyMo',
transformer_cfg=dict(
type='LLaMAHF',
block_size=1024, # max seq length
vocab_size=nb_code + 2, # append pad token and mask token
lm_head_cnt=vqvae_cfg['quantizer_cfg']['num_quantizers'],
**llama_configs[model_size],
),
vqvae_cfg=vqvae_cfg,
text_encoder_cfg = dict(
type='TextEncoder',
text_encoder_name = 'flan-t5-xl',
max_text_length=128,
output_dim=llama_configs[model_size]['n_embd'],
),
traj_encoder_cfg = dict(
type='Encoder',
input_emb_width=3,
down_t=2, norm=None, stride_t=2, width=512,
activation='relu', depth=3, dilation_growth_rate=3,
output_emb_width=llama_configs[model_size]['n_embd'],
),
audio_encoder_cfg=dict(
type='Encoder',
input_emb_width=512,
down_t=3, norm=None, stride_t=2, width=512,
activation='relu', depth=3, dilation_growth_rate=3,
output_emb_width=llama_configs[model_size]['n_embd'],
),
max_motion_len=300,
motion_downsample_ratio=unit_length,
audio_downsample_ratio=audio_unit_length,
stage=4,
cond_types=['text', 'traj']
# cond_types=['text']
)
# =============================================
# dataset
# =============================================
test_dataloader=dict(
batch_size=2,
num_workers=4,
pin_memory=True,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
collate_fn=dict(type='anymo_text_traj_collate_fn'),
dataset=dict(
type='AnyMoTextTrajDataset',
data_path_file='data/data_paths/omnihumo_motion_caption_dynamic_test.txt',
unit_length=unit_length,
mean_path=mean_path,
std_path=std_path,
traj_mean_path=traj_mean_path,
traj_std_path=traj_std_path,
min_motion_length=60,
max_motion_length=300,
# interval=10,
),
)
test_evaluator=dict(
# type='Text2MotionMetric',
type='Text2MotionTrajMetric',
tmr_cfg_path='work_dirs/temos_50e/temos.py',
tmr_ckpt_path='work_dirs/temos_50e/epoch_50.pth'
)
val_cfg = None
default_hooks=dict(
logger=dict(type='LoggerHook', interval=50),
)
|