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

train baseline

Browse Source
This commit is contained in:
luni
2023-07-02 15:00:33 +02:00
parent b387e9d64d
commit 455b5d02e7
4 changed files with 227 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/iccvw/abl_all.py
+195
View File
@@ -0,0 +1,195 @@
from utils.convnext_utils import CNBlockConfig
config = dict()
config['gpu'] = '2'
config['name'] = 'abl_iccvw_all_1'
config['notes'] = ''
config['tags'] = ['iccv', 'abl_all', 'torch']
config['script'] = 'scripts/cvae.py'
config['pretrained_model'] = ''
config['use_skinning'] = True # and deshaping
config['use_normals'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_norm_loss'] = True
ablation_count = [config['use_skinning'], config['use_normals'], config['use_extra_losses']].count(False)
assert ablation_count == 1 or ablation_count == 0
config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv
if config['use_skinning']:
assert (config['avg_garment_path'] != '')
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['batch_size'] = 4
config['resolution'] = 512
config['is_skirt'] = False
config['dataset_in_memory'] = False
config['is_cape'] = False
config['device'] = 'cuda'
config['train_dir'] = '/data-local/127/cloth3d_processed'
config['valid_dir'] = '/data-local/127/cloth3d_processed'
config['test_dir'] = '/data-net/dataset/test_ready'
if config['is_skirt']:
config['train_frames'] = 'frames_train_skirt.txt'
config['valid_frames'] = 'frames_valid_skirt.txt'
else:
config['train_frames'] = 'frames_train.txt'
config['valid_frames'] = 'frames_valid.txt'
config['static_model'] = 'checkpoints/iccvw/static_iccvw_4'
config['body_model'] = 'checkpoints/iccvw/body_iccvw_1'
config['normals_model'] = 'checkpoints/iccvw/norm_iccvw_1'
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['latent_dim'] = 128
###############
config['lr'] = 0.0001# 0.0001
###############
config['starting_balancing_coeff'] = 0.0001
config['balancing_activation_limit'] = 0.01
config['signature']: set = { 'id',
'frame_nr',
'pose_shape',
'uv_body_posed',
'mesh_body_posed',
'mesh_posed',
'uv_normals',
'uv_static',
'mesh_mask',
'uv_mask',
}
if config['use_skinning']:
config['signature'].add('uv_unposed')
config['signature'].add('mesh_unposed')
else:
config['signature'].add('uv_posed')
config['upscale_enhancer'] = ''
config['freeze_encoders'] = True
config['use_one_hot'] = False
config['bottleneck_scale'] = 2
config['cvae'] = dict()
config['cvae']['device'] = config['device']
config['cvae']['is_label'] = False
config['cvae']['input_features'] = 3+3+3+3 if config['use_normals'] else 3+3+3
config['cvae']['kernel_size'] = 7
config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale']
config['cvae']['latent_dim'] = config['cvae']['bottleneck_size']
block_setting = [
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 3),
CNBlockConfig(128, 256, 3),
CNBlockConfig(256, config['cvae']['bottleneck_size'], 9),
CNBlockConfig(config['cvae']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3),
CNBlockConfig(256, 128, 9),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 3),
CNBlockConfig(32, 3, 3),
]
config['cvae']['encoder_block_settings'] = block_setting
config['cvae']['decoder_block_settings'] = block_setting_reverse
####################### cond_ae_norm ############################
config['cond_ae_norm'] = dict()
config['cond_ae_norm']['device'] = config['device']
config['cond_ae_norm']['latent_dim'] = config['latent_dim']
config['cond_ae_norm']['kernel_size'] = 7
config['cond_ae_norm']['bottleneck_size'] = 40 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, 128, 9),
CNBlockConfig(128, config['cond_ae_norm']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_norm']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_norm']['encoder_block_settings'] = block_setting
config['cond_ae_norm']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_norm']['activation'] = 'tanh'
####################### cond_ae_body ############################
config['cond_ae_body'] = dict()
config['cond_ae_body']['device'] = config['device']
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['latent_dim'] = config['latent_dim']
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['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_body']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3),
CNBlockConfig(128, 64, 9),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 3, 3),
]
config['cond_ae_body']['encoder_block_settings'] = block_setting
config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_body']['activation'] = 'tanh'
####################### cond_ae_static ############################
config['cond_ae_static'] = dict()
config['cond_ae_static']['device'] = config['device']
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['latent_dim'] = config['latent_dim']
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
block_setting = [
CNBlockConfig(8, 16, 3),
CNBlockConfig(16, 32, 3),
CNBlockConfig(32, 64, 3),
CNBlockConfig(64, config['cond_ae_static']['bottleneck_size'], 9),
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3),
]
block_setting_reverse = [
CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3),
CNBlockConfig(64, 32, 9),
CNBlockConfig(32, 16, 3),
CNBlockConfig(16, 8, 3),
CNBlockConfig(8, 3, 3),
]
config['cond_ae_static']['encoder_block_settings'] = block_setting
config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse
config['cond_ae_static']['activation'] = 'tanh'
models/cvae_model.py
+6 -7
View File
@@ -18,14 +18,13 @@ class CVAE(nn.Module):
super(CVAE, self).__init__()
self.size = config['resolution']
if not('only_curvatures' in config.keys() and config['only_curvatures'] == True):
self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
if config['static_model'] != '':
load_model(self.static_enc, config['static_model'] + '_enc')
self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
if config['static_model'] != '':
load_model(self.static_enc, config['static_model'] + '_enc')
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
if config['body_model'] != '':
load_model(self.body_enc, config['body_model'] + '_enc')
if config['use_normals']:
self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device'])
scripts/cvae.py
+19 -8
View File
@@ -10,10 +10,12 @@ 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
if args['debug']:
from configs.abl_norm import config
from configs.iccvw.abl_all import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
@@ -42,6 +44,7 @@ from data.dataset import Dataset
from models.cond_encoder import AutoEncoder
from utils.checkpoints import CheckpointSaver, load_model
from utils.utils import InfiniteDataLoader as DataLoader
from utils.utils import text_to_html
from utils.postprocess import fix_collision
from losses.losses import collision_loss, normal_loss
from pytorch_msssim import MS_SSIM
@@ -51,7 +54,8 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/cvae.py')
artifact.add_file('losses/losses.py')
@@ -255,7 +259,14 @@ def plot_everything(batch: dict, type: str):
mean_hist = wandb.Histogram(sequence=mean.cpu().numpy().mean(axis=0))
logvar_hist = wandb.Histogram(sequence=logvar.cpu().numpy().mean(axis=0))
return { **figs, **figs_coll, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
batch_id = batch['id']
batch_frame_nr = batch['frame_nr'].cpu().numpy()
batch_garment = batch['garment']
batch_info = [f'batch_info: id: {batch_id[i]}, frame_nr: {batch_frame_nr[i]}, garment: {batch_garment[i]}' for i in range(len(batch_id))]
batch_info = text_to_html(batch_info)
batch_info = wandb.Html(batch_info)
return { **figs, **figs_coll, **{ f'{type}/batch_info': batch_info }, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
def sample(folder):
@@ -455,10 +466,10 @@ def train():
valid_accumulator.reset()
print(step)
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
if train_loss['train/euclid_error'] < config['balancing_activation_limit']:
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
if step % 25000 == 0:
ckpt_saver(model, step, optimizer, None)
step += 1
@@ -485,7 +496,7 @@ if __name__ == '__main__':
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
g_balancing_coeff = config['starting_balancing_coeff']
ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=10)
step = 0
if config['pretrained_model'] != '':
utils/utils.py
+7 -1
View File
@@ -42,4 +42,10 @@ class InfiniteDataLoader(DataLoader):
# Dataset exhausted, use a new fresh iterator.
self.dataset_iterator = super().__iter__()
batch = next(self.dataset_iterator)
return batch
return batch
def text_to_html(lines):
html = ''
for line in lines:
html += "<pre>" + line + "</pre>"
return html