This commit is contained in:
2022-09-02 14:15:54 +00:00
parent ee8f1cb68c
commit 977e49900d
16 changed files with 2241 additions and 272 deletions
+8 -8
View File
@@ -24,9 +24,9 @@ from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='ae_static_3'
name='ae_static_4'
def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
run = wandb.init(project='vae', name=name,
notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
@@ -389,13 +389,13 @@ def train():
batch = static_dataset_train.expand_batch(batch_to_dict(batch, signature_dynamic))
train_loss = train_step(batch)
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = valid_step(data_valid)
# data_valid = valid_iterator.get_next()
# data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = {}#valid_step(data_valid)
if step % 1000 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
valid_figs = {}#plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
@@ -432,8 +432,8 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/train', batch_size=8, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/train', signature=signature_static)
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=8, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=8, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)