Spaces:
Build error
Build error
File size: 3,219 Bytes
865fd8a |
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 |
import torch
import torch.nn as nn
import torch.optim as optim
class TrainWrapperBaseClass():
def __init__(self, args, config) -> None:
self.init_optimizer()
def init_optimizer(self) -> None:
print('using Adam')
self.generator_optimizer = optim.Adam(
self.generator.parameters(),
lr = self.config.Train.learning_rate.generator_learning_rate,
betas=[0.9, 0.999]
)
if self.discriminator is not None:
self.discriminator_optimizer = optim.Adam(
self.discriminator.parameters(),
lr = self.config.Train.learning_rate.discriminator_learning_rate,
betas=[0.9, 0.999]
)
def __call__(self, bat):
raise NotImplementedError
def get_loss(self, **kwargs):
raise NotImplementedError
def state_dict(self):
model_state = {
'generator': self.generator.state_dict(),
'generator_optim': self.generator_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 parameters(self):
return self.generator.parameters()
def load_state_dict(self, state_dict):
if 'generator' in state_dict:
self.generator.load_state_dict(state_dict['generator'])
else:
self.generator.load_state_dict(state_dict)
if 'generator_optim' in state_dict and self.generator_optimizer is not None:
self.generator_optimizer.load_state_dict(state_dict['generator_optim'])
if self.discriminator is not None:
self.discriminator.load_state_dict(state_dict['discriminator'])
if 'discriminator_optim' in state_dict and self.discriminator_optimizer is not None:
self.discriminator_optimizer.load_state_dict(state_dict['discriminator_optim'])
def infer_on_audio(self, aud_fn, initial_pose=None, norm_stats=None, **kwargs):
raise NotImplementedError
def init_params(self):
if self.config.Data.pose.convert_to_6d:
scale = 2
else:
scale = 1
global_orient = round(0 * scale)
leye_pose = reye_pose = round(0 * scale)
jaw_pose = round(0 * scale)
body_pose = round((63 - 24) * scale)
left_hand_pose = right_hand_pose = round(45 * scale)
if self.expression:
expression = 100
else:
expression = 0
b_j = 0
jaw_dim = jaw_pose
b_e = b_j + jaw_dim
eye_dim = leye_pose + reye_pose
b_b = b_e + eye_dim
body_dim = global_orient + body_pose
b_h = b_b + body_dim
hand_dim = left_hand_pose + right_hand_pose
b_f = b_h + hand_dim
face_dim = expression
self.dim_list = [b_j, b_e, b_b, b_h, b_f]
self.full_dim = jaw_dim + eye_dim + body_dim + hand_dim
self.pose = int(self.full_dim / round(3 * scale))
self.each_dim = [jaw_dim, eye_dim + body_dim, hand_dim, face_dim] |