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
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'] != '':