import torch
import config
from loss import linearizedlie_loss
import utils
from choose_dataset import DatasetChooser
from torch import autograd
from STLN import ST_HMR, ST_LSTM
import choose_dataset
from torch.utils.data import DataLoader
if __name__ == '__main__':
config = config.TrainConfig('Human', 'lie', 'all')
# choose = DatasetChooser(config)
# prediction_dataset, _ = choose(prediction=True)
# prediction_loader = DataLoader(prediction_dataset, batch_size=config.batch_size, shuffle=True)
choose = DatasetChooser(config)
train_dataset, bone_length = choose(train=True)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=True)
test_dataset, _ = choose(train=False)
test_loader = DataLoader(test_dataset,
batch_size=config.batch_size,
shuffle=True)
prediction_dataset, bone_length = choose(prediction=True)
prediction_loader = DataLoader(prediction_dataset,
batch_size=config.batch_size,
shuffle=True)
def prediction(config, checkpoint_dir, output_dir):
print('Start testing the Model!')
if not (os.path.exists(output_dir)):
os.makedirs(output_dir)
print("Outputs saved to: " + output_dir)
# generate data loader,更改预测时输出的视频帧长
if config.dataset == 'Mouse':
config.output_window_size = 7
else:
config.output_window_size = 50
choose = DatasetChooser(config)
if config.dataset == 'Human':
# This step is to get mean value, etc for unnorm
_, _ = choose(train=True)
if config.longterm is False:
prediction_dataset, bone_length = choose(prediction=True)
x_test, y_test, dec_in_test = prediction_dataset
actions = list(x_test.keys())
else:
# get raw validation data because the test data isn't usable
train_dataset, bone_length = choose(train=False)
test_set = train_dataset.data
x_test = {}
y_test = {}
dec_in_test = {}
test_set = test_set[0]
x_test[config.filename] = np.reshape(
test_set[:config.input_window_size - 1, :],
[1, -1, config.input_size])
y_test[config.filename] = np.reshape(
test_set[config.input_window_size:, :], [1, -1, config.input_size])
dec_in_test[config.filename] = np.reshape(
test_set[config.input_window_size - 1:-1, :],
[1, -1, config.input_size])
config.output_window_size = y_test[config.filename].shape[1]
actions = [config.filename]
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('Device {} will be used to save parameters'.format(device))
net = choose_net(config)
net.to(device)
print('Total param number:' + str(sum(p.numel()
for p in net.parameters())))
print('Encoder param number:' +
str(sum(p.numel() for p in net.encoder_cell.parameters())))
print('Decoder param number:' +
str(sum(p.numel() for p in net.decoder.parameters())))
net = torch.nn.DataParallel(net)
dir = utils.get_file_list(checkpoint_dir)
print('Load model from:' + checkpoint_dir + dir[-1]) # 使用最后保存的checkpoint
## 用GPU时记得改成 map_location='cuda:0',用cpu时是 map_location=torch.device('cpu')
net.load_state_dict(
torch.load(checkpoint_dir + dir[-1], map_location='cuda:0'))
y_predict = {}
with torch.no_grad():
for act in actions:
x_test_ = torch.from_numpy(x_test[act]).float().to(device)
dec_in_test_ = torch.from_numpy(
dec_in_test[act]).float().to(device)
pred = net(x_test_, dec_in_test_, train=False)
pred = pred.cpu().numpy()
y_predict[act] = pred
for act in actions:
if config.datatype == 'smpl':
mean_error, _ = utils.mean_euler_error(config, act, y_predict[act],
y_test[act])
sio.savemat(output_dir + 'error_' + act + '.mat',
dict([('error', mean_error)]))
for i in range(y_predict[act].shape[0]):
if config.dataset == 'Human':
y_p = utils.unNormalizeData(y_predict[act][i],
config.data_mean,
config.data_std,
config.dim_to_ignore)
y_t = utils.unNormalizeData(y_test[act][i],
config.data_mean,
config.data_std,
config.dim_to_ignore)
# expmap_all = utils.revert_coordinate_space(np.vstack((y_t, y_p)), np.eye(3), np.zeros(3))
# y_p = expmap_all[config.output_window_size:]
# y_t = expmap_all[:config.output_window_size]
# else:
# y_p = y_predict[act][i]
# y_t = y_test[act][i]
# y_p_xyz = utils.fk(y_p, config, bone_length)
# y_t_xyz = utils.fk(y_t, config, bone_length)
sio.savemat(
output_dir + 'prediction_smpl_' + act + '_' + str(i) +
'.mat', dict([('prediction', y_p)]))
sio.savemat(
output_dir + 'gt_smpl_' + act + '_' + str(i) + '.mat',
dict([('gt', y_t)]))
def train(config, checkpoint_dir):
print('Start Training the Model!')
# generate data loader
if config.longterm is True:
config.output_window_size = 100
# from choose_dataset import DatasetChooser
choose = DatasetChooser(config)
train_dataset, bone_length = choose(train=True)
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=True)
test_dataset, _ = choose(train=False)
test_loader = DataLoader(test_dataset,
batch_size=config.batch_size,
shuffle=True)
prediction_dataset, bone_length = choose(prediction=True)
x_test, y_test, dec_in_test = prediction_dataset
device = torch.device(
"cuda:" +
str(config.device_ids[0]) if torch.cuda.is_available() else "cpu")
print('Device {} will be used to save parameters'.format(device))
torch.cuda.manual_seed_all(112858)
net = choose_net(config)
net.to(device)
print('Total param number:' + str(sum(p.numel()
for p in net.parameters())))
print('Encoder param number:' +
str(sum(p.numel() for p in net.encoder_cell.parameters())))
print('Decoder param number:' +
str(sum(p.numel() for p in net.decoder.parameters())))
if torch.cuda.device_count() > 1:
print("{} GPUs are usable!".format(str(torch.cuda.device_count())))
net = torch.nn.DataParallel(net, device_ids=config.device_ids)
if config.restore is True:
dir = utils.get_file_list(checkpoint_dir)
print('Load model from:' + checkpoint_dir + dir[-1])
net.load_state_dict(
torch.load(checkpoint_dir + dir[-1], map_location='cuda:0'))
optimizer = optim.Adam(net.parameters(), lr=config.learning_rate)
# Save model
if not (os.path.exists(checkpoint_dir)):
os.makedirs(checkpoint_dir)
best_error = float('inf')
best_error_list = None
# log 定义两个数组
trainloss_list = []
validloss_list = []
Error_list = []
for epoch in range(config.max_epoch):
print("At epoch:{}".format(str(epoch + 1)))
prog = Progbar(target=config.training_size)
prog_valid = Progbar(target=config.validation_size)
# Train
#with torch.autograd.set_detect_anomaly(True):
for it in range(config.training_size):
for i, data in enumerate(train_loader, 0):
encoder_inputs = data['encoder_inputs'].float().to(
device) ## 前t-1帧
decoder_inputs = data['decoder_inputs'].float().to(
device) ## t-1到t-1+output_window_size帧
decoder_outputs = data['decoder_outputs'].float().to(
device) ## t帧到以后的
prediction = net(encoder_inputs, decoder_inputs, train=True)
loss = Loss(prediction, decoder_outputs, bone_length, config)
net.zero_grad()
loss.backward()
_ = torch.nn.utils.clip_grad_norm_(net.parameters(), 5)
optimizer.step()
prog.update(it + 1, [("Training Loss", loss.item())])
trainloss_list.append(loss.item)
# valid
with torch.no_grad():
for it in range(config.validation_size):
for j in range(3):
for i, data in enumerate(test_loader, 0):
if j == 0 and i == 0:
encoder_inputs = data['encoder_inputs'].float().to(
device)
decoder_inputs = data['decoder_inputs'].float().to(
device)
decoder_outputs = data['decoder_outputs'].float(
).to(device)
else:
encoder_inputs = torch.cat([
data['encoder_inputs'].float().to(device),
encoder_inputs
],
dim=0)
decoder_inputs = torch.cat([
data['decoder_inputs'].float().to(device),
decoder_inputs
],
dim=0)
decoder_outputs = torch.cat([
data['decoder_outputs'].float().to(device),
decoder_outputs
],
dim=0)
prediction = net(encoder_inputs, decoder_inputs, train=True)
loss = Loss(prediction, decoder_outputs, bone_length, config)
prog_valid.update(it + 1, [("Testing Loss", loss.item())])
validloss_list.append(loss.item)
#Test prediction
actions = list(x_test.keys())
y_predict = {}
with torch.no_grad():
for act in actions:
x_test_ = torch.from_numpy(x_test[act]).float().to(device)
dec_in_test_ = torch.from_numpy(
dec_in_test[act]).float().to(device)
pred = net(x_test_, dec_in_test_, train=False)
pred = pred.cpu().numpy()
y_predict[act] = pred
error_actions = 0.0
for act in actions:
if config.datatype == 'lie':
mean_error, _ = utils.mean_euler_error(
config, act, y_predict[act],
y_test[act][:, :config.output_window_size, :])
error = mean_error[[1, 3, 7, 9]]
if config.datatype == 'smpl':
mean_error, _ = utils.mean_euler_error(
config, act, y_predict[act],
y_test[act][:, :config.output_window_size, :])
error = mean_error[[1, 3, 7, 9]]
# error_actions += error.mean()
error_actions += mean_error.mean()
error_actions /= len(actions)
Error_list.append(error_actions)
if error_actions < best_error:
print(error_actions)
print(best_error)
best_error_list = error
best_error = error_actions
torch.save(net.state_dict(),
checkpoint_dir + 'Epoch_' + str(epoch + 1) + '.pth')
print('Current best:' + str(round(best_error_list[0], 2)) + ' ' +
str(round(best_error_list[1], 2)) + ' ' +
str(round(best_error_list[2], 2)) + ' ' +
str(round(best_error_list[3], 2)))
if not (os.path.exists('./log')):
os.makedirs('./log')
time_now = datetime.datetime.now().strftime("%Y_%m_%d_%H_%M_%S")
np.savez('./log/' + time_now + '.npz',
trainloss_list=trainloss_list,
validloss_list=validloss_list,
Error_list=Error_list)
def predict_one(config, input_pose, checkpoint_dir):
# 更改预测时输出的视频帧长
config.output_window_size = 90
# 指定抽样后的起始帧 1: 470
start_frame = 330
datetime_p = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
print('Start predict at time: ', datetime_p)
output_dir = './predict_one_output/pose_' + datetime_p + '.npz'
error_dir = './predict_one_output/error_' + datetime_p + '.mat'
if not (os.path.exists('./predict_one_output/')):
os.makedirs('./predict_one_output/')
print("pose paras will be saved to: " + output_dir)
print("errors will be saved to: " + error_dir)
# This step is to get mean value,dim_to_use etc for norm and unnorm
choose = DatasetChooser(config)
_, _ = choose(train=True)
x_test = {}
y_test = {}
dec_in_test = {}
batch_size = 1
source_seq_len = config.input_window_size
target_seq_len = config.output_window_size
encoder_inputs = np.zeros((batch_size, source_seq_len - 1, 60),
dtype=float)
decoder_inputs = np.zeros((batch_size, target_seq_len, 60), dtype=float)
decoder_outputs = np.zeros((batch_size, target_seq_len, 60), dtype=float)
data_seq = input_pose[start_frame:start_frame + source_seq_len +
target_seq_len, :66] # 70（20+50）帧，每帧66个pose参数
test_data = normalize_data(data_seq, config.data_mean, config.data_std,
config.dim_to_use)
for i in range(batch_size):
encoder_inputs[i, :, :] = test_data[0:source_seq_len - 1][:] # x_test
decoder_inputs[i, :, :] = test_data[source_seq_len -
1:(source_seq_len +
target_seq_len -
1)][:] # decoder_in_test
decoder_outputs[i, :, :] = test_data[
source_seq_len:][:] # y_test, 相比于decoder_inputs 后移一位
x_test['walking'] = encoder_inputs
y_test['walking'] = decoder_outputs
dec_in_test['walking'] = np.zeros(
[decoder_inputs.shape[0], 1, decoder_inputs.shape[2]])
dec_in_test['walking'][:, 0, :] = decoder_inputs[:, 0, :]
actions = list(x_test.keys()) # ['walking']
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('Device {} will be used to save parameters'.format(device))
net = ST_HRN(config)
net.to(device)
print('Total param number:' + str(sum(p.numel()
for p in net.parameters())))
print('Encoder param number:' +
str(sum(p.numel() for p in net.encoder_cell.parameters())))
print('Decoder param number:' +
str(sum(p.numel() for p in net.decoder.parameters())))
net = torch.nn.DataParallel(net)
# dir = utils.get_file_list(checkpoint_dir)
# print('Load model from:' + checkpoint_dir + dir[-1]) # 使用最后保存的checkpoint
## 用GPU时记得改成 map_location='cuda:0',用cpu时是 map_location=torch.device('cpu')
# net.load_state_dict(torch.load(checkpoint_dir + dir[-1], map_location='cuda:0'))
net.load_state_dict(torch.load(checkpoint_dir, map_location='cuda:0'))
y_predict = {}
with torch.no_grad():
for act in actions:
x_test_ = torch.from_numpy(x_test[act]).float().to(device)
dec_in_test_ = torch.from_numpy(
dec_in_test[act]).float().to(device)
pred = net(x_test_, dec_in_test_, train=False)
pred = pred.cpu().numpy()
y_predict[act] = pred
for act in actions:
if config.datatype == 'smpl':
mean_error, _ = utils.mean_euler_error(config, act, y_predict[act],
y_test[act])
sio.savemat(error_dir, dict([('error', mean_error)]))
for i in range(y_predict[act].shape[0]):
print(y_predict[act].shape[0])
if config.dataset == 'Human':
y_p = utils.unNormalizeData(y_predict[act][i],
config.data_mean,
config.data_std,
config.dim_to_ignore)
y_t = utils.unNormalizeData(y_test[act][i],
config.data_mean,
config.data_std,
config.dim_to_ignore)
# 保存
np.savez(output_dir, y_p=y_p, y_t=y_t)