-
Notifications
You must be signed in to change notification settings - Fork 0
/
train.py
103 lines (80 loc) · 3.88 KB
/
train.py
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
import os
import sys
import pickle
import argparse
import os.path as osp
import torch
from tqdm import tqdm
sys.path.append(os.getcwd())
from src.utils.general import set_random_seed
from src.utils.euler import define_actions
from src.config import Config
from src.net import Parallel_Denoiser, Series_Denoiser
from src.diff import DDPM
from src.train.euler import train_euler
from src.train.xyz import train_xyz
from src.dataset.xyz import DatasetH36M, DatasetHumanEva
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', default='h36m_euler_series_20step')
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--gpu-id', type=int, default=0)
args = parser.parse_args()
cfg = Config(args.cfg)
set_random_seed(args.seed)
device = 'cuda:{}'.format(args.gpu_id) if args.gpu_id >= 0 else 'cpu'
os.makedirs('./log', exist_ok=True)
f = open('./log/[TRAIN]{}.txt'.format(args.cfg), 'w')
if 'euler' in args.cfg:
# fixed prefix and pred length
prefix_len = 50
pred_len = 25
pose_dim = 54
actions = define_actions("all")
data = pickle.load(open('./data/h36m_euler.pkl', 'rb'))
train_set = data['train']
test_set = data['test']
data_mean = data['mean']
data_std = data['std']
dim_to_ignore = data['dim_to_ignore']
dim_to_use = data['dim_to_use']
## define network
if '_parallel_' in args.cfg: # "parallel" model in paper
denoiser = Parallel_Denoiser(pose_dim, cfg.qkv_dim, cfg.num_layers, cfg.num_heads,
prefix_len, pred_len, cfg.diff_steps)
elif '_series_' in args.cfg: # "series" model in paper.
denoiser = Series_Denoiser(pose_dim, cfg.qkv_dim, cfg.num_layers, cfg.num_heads,
prefix_len, pred_len, cfg.diff_steps)
ddpm = DDPM(denoiser, cfg, device).to(device)
optimizer = torch.optim.Adam(ddpm.parameters(), lr=cfg.learning_rate)
for epoch in tqdm(range(cfg.max_epoch)):
train_euler(f, epoch, cfg.epoch_iters, ddpm, optimizer, train_set, cfg.batch_size, pose_dim,
prefix_len, pred_len, actions, device)
if (epoch+1)%cfg.save_epoch == 0:
torch.save(ddpm.state_dict(), osp.join(cfg.model_dir, '{:04}.pth'.format(epoch+1)))
torch.save(ddpm.state_dict(), osp.join(cfg.model_dir, 'recent.pth'))
f.close()
elif 'xyz' in args.cfg:
# fixed values! do not change.
prefix_len = 25
pred_len = 100
if 'h36m' in args.cfg:
dataset_cls = DatasetH36M
elif 'humaneva' in args.cfg:
dataset_cls = DatasetHumanEva
dataset = dataset_cls('train', prefix_len, pred_len, actions='all')
## define network
if '_parallel_' in args.cfg: # "parallel" model in paper
denoiser = Parallel_Denoiser(dataset.traj_dim, cfg.qkv_dim, cfg.num_layers, cfg.num_heads,
prefix_len, pred_len, cfg.diff_steps)
elif '_series_' in args.cfg: # "series" model in paper
denoiser = Series_Denoiser(dataset.traj_dim, cfg.qkv_dim, cfg.num_layers, cfg.num_heads,
prefix_len, pred_len, cfg.diff_steps)
ddpm = DDPM(denoiser, cfg, device).to(device)
optimizer = torch.optim.Adam(ddpm.parameters(), lr=cfg.learning_rate)
for epoch in tqdm(range(cfg.max_epoch)):
train_xyz(f, epoch, cfg.epoch_iters, ddpm, optimizer, dataset, cfg.batch_size, prefix_len, device)
if (epoch+1)%cfg.save_epoch == 0:
torch.save(ddpm.state_dict(), osp.join(cfg.model_dir, '{:04}.pth'.format(epoch+1)))
torch.save(ddpm.state_dict(), osp.join(cfg.model_dir, 'recent.pth'))
f.close()