DisMouse/main.py
Guanhua Zhang b102c2a534 init
2024-10-08 16:23:07 +02:00

29 lines
No EOL
780 B
Python

import os
from templates import *
if __name__ == '__main__':
gpus = [0]
conf = mouse_autoenc('trainDiff')
betas = {
'recon':1,
'noise':1,
'user':0.01,
'nonuser':0.01,
'mi': 0.01
}
betastr = ''
for k,v in betas.items():
betastr += f'{k}{v}_'
betastr = betastr[:-1]
diffmodel = LitModel(conf, betas)
conf.logdir= f'{conf.logdir}/mouse_autoenc/{conf.pretrainDataset}/{betastr}/embDim{conf.net_beatgans_embed_channels}/win{conf.AEwin}/slid{conf.slid}/GRL/'
MakeDir(conf.logdir)
os.environ['WANDB_CACHE_DIR'] = conf.logdir
os.environ['WANDB_DATA_DIR'] = conf.logdir
os.environ['WANDB_IGNORE_GLOBS'] = '*.ckpt'
train(conf, diffmodel, gpus=gpus)