File size: 1,796 Bytes
8e5456b | 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 | #!/usr/bin/env python3
"""Save an UNTRAINED stage-2 checkpoint, as the iter-0 control for eval_vsl.py.
Without this baseline, a flat generation-metric curve is ambiguous: it could mean
"training converged in the first few thousand iters" or "the transformer barely
contributes and the number reflects the pose prior". Scoring a randomly
initialized GPT through the *trained* VQ-VAE decoder separates the two.
Reuses the trained checkpoint's args so the architecture and eval path are
identical -- only the transformer weights are random.
"""
import argparse
import torch
import models.t2m_trans as trans
from models.text_encoder_vi import ViTextEncoder
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--like', required=True, help='trained stage-2 ckpt to copy args from')
ap.add_argument('--out', required=True)
ap.add_argument('--seed', type=int, default=1234)
args = ap.parse_args()
ref = torch.load(args.like, map_location='cpu')
gargs = argparse.Namespace(**ref['args'])
torch.manual_seed(args.seed)
# text encoder is only needed for its hidden size (it is frozen either way)
dim = ViTextEncoder(gargs.text_model, device='cpu').dim
net = trans.Text2Motion_Transformer(
num_vq=gargs.nb_code, embed_dim=gargs.embed_dim_gpt, clip_dim=dim,
block_size=gargs.max_tokens + 1, num_layers=gargs.num_layers,
n_head=gargs.n_head_gpt, drop_out_rate=gargs.drop_out_rate, fc_rate=gargs.ff_rate)
torch.save({'trans': net.state_dict(), 'args': ref['args'], 'iter': 0,
'val_loss': float('nan'), 'val_acc': float('nan'),
'note': 'UNTRAINED random init - iter-0 control'}, args.out)
print(f'wrote untrained control checkpoint -> {args.out}')
if __name__ == '__main__':
main()
|