update

example/config.json

@@ -7,8 +7,8 @@
   "out_dir": "./results",
   "device": "auto",
   "timesteps": 600,
-  "batch_size": 24,
-  "seq_len": 128,
+  "batch_size": 12,
+  "seq_len": 96,
   "epochs": 10,
   "max_batches": 4000,
   "lambda": 0.7,
@@ -26,7 +26,7 @@
   "use_tanh_eps": false,
   "eps_scale": 1.0,
   "model_time_dim": 128,
-  "model_hidden_dim": 640,
+  "model_hidden_dim": 512,
   "model_num_layers": 2,
   "model_dropout": 0.1,
   "model_ff_mult": 2,
@@ -35,7 +35,7 @@
   "backbone_type": "transformer",
   "transformer_num_layers": 3,
   "transformer_nhead": 4,
-  "transformer_ff_dim": 768,
+  "transformer_ff_dim": 512,
   "transformer_dropout": 0.1,
   "disc_mask_scale": 0.9,
   "cont_loss_weighting": "inv_std",
@@ -58,14 +58,14 @@
   "shuffle_buffer": 256,
   "use_temporal_stage1": true,
   "temporal_backbone": "transformer",
-  "temporal_hidden_dim": 384,
+  "temporal_hidden_dim": 256,
   "temporal_num_layers": 1,
   "temporal_dropout": 0.0,
   "temporal_pos_dim": 64,
   "temporal_use_pos_embed": true,
   "temporal_transformer_num_layers": 2,
   "temporal_transformer_nhead": 4,
-  "temporal_transformer_ff_dim": 512,
+  "temporal_transformer_ff_dim": 256,
   "temporal_transformer_dropout": 0.1,
   "temporal_epochs": 3,
   "temporal_lr": 0.001,
@@ -74,6 +74,6 @@
   "snr_weighted_loss": true,
   "snr_gamma": 1.0,
   "residual_stat_weight": 0.05,
-  "sample_batch_size": 8,
-  "sample_seq_len": 128
+  "sample_batch_size": 4,
+  "sample_seq_len": 96
 }