Spaces:
Running
Running
| import os | |
| import sys | |
| sys.path.append(os.getcwd()) | |
| from nets.base import TrainWrapperBaseClass | |
| from nets.spg.s2glayers import Discriminator as D_S2G | |
| from nets.spg.vqvae_1d import AE as s2g_body | |
| import torch | |
| import torch.optim as optim | |
| import torch.nn.functional as F | |
| from data_utils.lower_body import c_index, c_index_3d, c_index_6d | |
| def separate_aa(aa): | |
| aa = aa[:, :, :].reshape(aa.shape[0], aa.shape[1], -1, 5) | |
| axis = F.normalize(aa[:, :, :, :3], dim=-1) | |
| angle = F.normalize(aa[:, :, :, 3:5], dim=-1) | |
| return axis, angle | |
| class TrainWrapper(TrainWrapperBaseClass): | |
| ''' | |
| a wrapper receving a batch from data_utils and calculate loss | |
| ''' | |
| def __init__(self, args, config): | |
| self.args = args | |
| self.config = config | |
| self.device = torch.device(self.args.gpu) | |
| self.global_step = 0 | |
| self.gan = False | |
| self.convert_to_6d = self.config.Data.pose.convert_to_6d | |
| self.preleng = self.config.Data.pose.pre_pose_length | |
| self.expression = self.config.Data.pose.expression | |
| self.epoch = 0 | |
| self.init_params() | |
| self.num_classes = 4 | |
| self.g = s2g_body(self.each_dim[1] + self.each_dim[2], embedding_dim=64, num_embeddings=0, | |
| num_hiddens=1024, num_residual_layers=2, num_residual_hiddens=512).to(self.device) | |
| if self.gan: | |
| self.discriminator = D_S2G( | |
| pose_dim=110 + 64, pose=self.pose | |
| ).to(self.device) | |
| else: | |
| self.discriminator = None | |
| if self.convert_to_6d: | |
| self.c_index = c_index_6d | |
| else: | |
| self.c_index = c_index_3d | |
| super().__init__(args, config) | |
| def init_optimizer(self): | |
| self.g_optimizer = optim.Adam( | |
| self.g.parameters(), | |
| lr=self.config.Train.learning_rate.generator_learning_rate, | |
| betas=[0.9, 0.999] | |
| ) | |
| def state_dict(self): | |
| model_state = { | |
| 'g': self.g.state_dict(), | |
| 'g_optim': self.g_optimizer.state_dict(), | |
| 'discriminator': self.discriminator.state_dict() if self.discriminator is not None else None, | |
| 'discriminator_optim': self.discriminator_optimizer.state_dict() if self.discriminator is not None else None | |
| } | |
| return model_state | |
| def __call__(self, bat): | |
| # assert (not self.args.infer), "infer mode" | |
| self.global_step += 1 | |
| total_loss = None | |
| loss_dict = {} | |
| aud, poses = bat['aud_feat'].to(self.device).to(torch.float32), bat['poses'].to(self.device).to(torch.float32) | |
| # id = bat['speaker'].to(self.device) - 20 | |
| # id = F.one_hot(id, self.num_classes) | |
| poses = poses[:, self.c_index, :] | |
| gt_poses = poses[:, :, self.preleng:].permute(0, 2, 1) | |
| loss = 0 | |
| loss_dict, loss = self.vq_train(gt_poses[:, :], 'g', self.g, loss_dict, loss) | |
| return total_loss, loss_dict | |
| def vq_train(self, gt, name, model, dict, total_loss, pre=None): | |
| x_recon = model(gt_poses=gt, pre_state=pre) | |
| loss, loss_dict = self.get_loss(pred_poses=x_recon, gt_poses=gt, pre=pre) | |
| # total_loss = total_loss + loss | |
| if name == 'g': | |
| optimizer_name = 'g_optimizer' | |
| optimizer = getattr(self, optimizer_name) | |
| optimizer.zero_grad() | |
| loss.backward() | |
| optimizer.step() | |
| for key in list(loss_dict.keys()): | |
| dict[name + key] = loss_dict.get(key, 0).item() | |
| return dict, total_loss | |
| def get_loss(self, | |
| pred_poses, | |
| gt_poses, | |
| pre=None | |
| ): | |
| loss_dict = {} | |
| rec_loss = torch.mean(torch.abs(pred_poses - gt_poses)) | |
| v_pr = pred_poses[:, 1:] - pred_poses[:, :-1] | |
| v_gt = gt_poses[:, 1:] - gt_poses[:, :-1] | |
| velocity_loss = torch.mean(torch.abs(v_pr - v_gt)) | |
| if pre is None: | |
| f0_vel = 0 | |
| else: | |
| v0_pr = pred_poses[:, 0] - pre[:, -1] | |
| v0_gt = gt_poses[:, 0] - pre[:, -1] | |
| f0_vel = torch.mean(torch.abs(v0_pr - v0_gt)) | |
| gen_loss = rec_loss + velocity_loss + f0_vel | |
| loss_dict['rec_loss'] = rec_loss | |
| loss_dict['velocity_loss'] = velocity_loss | |
| # loss_dict['e_q_loss'] = e_q_loss | |
| if pre is not None: | |
| loss_dict['f0_vel'] = f0_vel | |
| return gen_loss, loss_dict | |
| def load_state_dict(self, state_dict): | |
| self.g.load_state_dict(state_dict['g']) | |
| def extract(self, x): | |
| self.g.eval() | |
| if x.shape[2] > self.full_dim: | |
| if x.shape[2] == 239: | |
| x = x[:, :, 102:] | |
| x = x[:, :, self.c_index] | |
| feat = self.g.encode(x) | |
| return feat.transpose(1, 2), x | |