luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

adapt ae_body to cape

Browse Source
This commit is contained in:
luni
2023-05-21 19:45:24 +02:00
parent 75df8a09cc
commit 70162cf8ac
3 changed files with 108 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/ae_body_cape_m.py
+73
View File
@@ -0,0 +1,73 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '4'
config['name'] = 'body_torch_cape_m'
config['notes'] = ' '
config['tags'] = ['torch', 'body']
config['script'] = 'scripts/ae_body_train.py'
config['pretrained_model'] = ''
config['batch_size'] = 4
config['train_dir'] = '/data-local/127/cape_processed/train/male'
config['valid_dir'] = '/data-local/127/cape_processed/test/male'
config['test_dir'] = '/data-local/127/cape_processed/test/male'
config['train_frames'] = None
config['valid_frames'] = None
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = True
config['remove_global_orientation'] = True
config['device'] = 'cuda'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
config['bottleneck_scale'] = 2
config['cond_ae'] = dict()
config['cond_ae']['device'] = config['device']
config['cond_ae']['is_label'] = False
config['cond_ae']['latent_dim'] = config['latent_dim']
config['cond_ae']['kernel_size'] = 7
config['cond_ae']['bottleneck_size'] = 16 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
# CNBlockConfig(3, 3, 3),
]
config['cond_ae']['encoder_block_settings'] = block_setting
config['cond_ae']['decoder_block_settings'] = block_setting_reverse
config['cond_ae']['activation'] = 'tanh'
config['lr'] = 0.0001
config['cond_ae']['upscale_output'] = False
config['signature'] = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_mask',
'mesh_body_posed',
}
losses/losses.py
+18 -8
View File
@@ -24,14 +24,14 @@ F_norm_skirt = torch.tensor(np.tile(np.load("smpl/smpl_faces_skirt.npy")[to_keep
F_norm_orig = torch.tensor(np.tile(np.load("smpl/smpl_faces_orig.npy")[:6890,:][None,:,:], [4,1,1]), device='cuda')
def _nearest_neighbour(V, B):
def _nearest_neighbour(V: torch.Tensor, B: torch.Tensor):
# flatten V and B along the batch dimension, create a batch_x and batch_y variable to restore the batch dimension after the nearest neighbour search using pytorch_geomtry knn function
batch_size = V.shape[0]
batch_x = torch.arange(batch_size, device=V.get_device()).repeat_interleave(V.shape[1])
batch_y = torch.arange(batch_size, device=B.get_device()).repeat_interleave(B.shape[1])
V = V.reshape(-1, V.shape[-1])
B = B.reshape(-1, B.shape[-1])
idx = knn(V, B, k=1, batch_x=batch_x, batch_y=batch_y)[1]
idx = knn(V, B.to(V.dtype), k=1, batch_x=batch_x, batch_y=batch_y)[1]
idx = (idx - batch_x * 14475).reshape(batch_size, 14475)
return idx
@@ -64,18 +64,28 @@ edges: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())
laplacian_loss_skirt = sio.loadmat("custom_utils/ssmpl_skirt0.mat")['L'].tocoo()
edges_skirt: torch.Tensor = torch.tensor(np.stack(np.nonzero(laplacian_loss.todense())).transpose(), device='cuda')
smpl_to_keep = sio.loadmat('smpl/keep_remove.mat')['keep'].reshape(-1)
smpl_to_keep = smpl_to_keep[smpl_to_keep < 6890]
laplacian_loss_orig = laplacian_loss.todense()[:6890,:6890][smpl_to_keep,:][:,smpl_to_keep]
edges_cape = torch.tensor(np.stack(np.nonzero(laplacian_loss_orig)).transpose(), device='cuda')
def get_edge_mask(masks: torch.Tensor):
def get_edge_mask(edges: torch.Tensor, masks: torch.Tensor):
if edges.get_device() != masks.get_device():
edges.to(masks.get_device())
return torch.min(masks[:,edges], dim=-1)[0]
def edge_loss(x, y, mask):
e_mask = get_edge_mask(mask)
x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
if x.shape[1] == 14475:
edges_local = edges
elif x.shape[1] == 3638:
edges_local = edges_cape
else:
raise ValueError("Invalid number of vertices")
e_mask = get_edge_mask(edges_local, mask)
x_e = x[:,edges_local[:,0],:] - x[:,edges_local[:,1],:]
x_e = torch.sum(x_e ** 2, -1)
x_e = torch.sqrt(x_e + 1e-9)
y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
y_e = y[:,edges_local[:,0],:] - y[:,edges_local[:,1],:]
y_e = torch.sum(y_e ** 2, -1)
y_e = torch.sqrt(y_e + 1e-9)
d_e = y_e - x_e
@@ -173,8 +183,8 @@ def log_normal_pdf(sample, mean, logvar, raxis=1):
def kl_loss(mean, logvar):
kl_loss = 1 + logvar - tf.square(mean) - tf.exp(logvar)
kl_loss = tf.reduce_mean(kl_loss)
kl_loss = 1 + logvar - torch.square(mean) - torch.exp(logvar)
kl_loss = torch.mean(kl_loss)
kl_loss *= -0.5
return kl_loss
scripts/ae_body_train.py
+17 -16
View File
@@ -8,14 +8,16 @@ import importlib
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
parser.add_argument('--config', type=str)
parser.add_argument('--debug', action='store_true')
args = vars(parser.parse_args())
config_module = importlib.import_module(args['config'])
config = config_module.config
# from configs.ae_body_1 import config
# from configs.ae_body_cape_m import config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
import time
import wandb
@@ -155,37 +157,31 @@ def train():
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
wandb.log(to_log, step=step)
if not args['debug']:
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
if step % 10000 == 0:
ckpt_saver_enc(model.encoder, step, optimizer)
ckpt_saver_dec(model.decoder, step, optimizer)
ckpt_saver_enc(model.encoder, step, optimizer, scaler=None)
ckpt_saver_dec(model.decoder, step, optimizer, scaler=None)
step += 1
if __name__ == '__main__':
# init_wandb()
# device = "cuda" if torch.cuda.is_available() else "cpu"
if not args['debug']:
init_wandb()
# parser = argparse.ArgumentParser(description='')
# parser.add_argument('--name', type=str)
# args = vars(parser.parse_args())
resolution_modifier = ''
if config['resolution'] != 512:
resolution_modifier = '_' + str(config['resolution'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True)
valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
train_accumulator = Accumulator()
@@ -206,6 +202,11 @@ if __name__ == '__main__':
run_name = config['name']
print(f'started training: {run_name}')
if not args['debug']:
# model = torch.compile(model)
wandb.watch(model, log='all', log_freq=1000, log_graph=True)
train()