speeeeed

train.py

@@ -27,19 +27,19 @@ import torch.nn.functional as F
 # --------------------------- Параметры ---------------------------
 ds_path = "datasets/384"
 project = "2b"
-batch_size = 200 #50
+batch_size = 16 #50
 base_learning_rate = 8e-5
 min_learning_rate = 4e-5
-num_epochs = 40
+num_epochs = 10
 # samples/save per epoch
-sample_interval_share = 5
+sample_interval_share = 10
 use_wandb = True
 save_model = True
 use_decay = True
 fbp = False # fused backward pass
 optimizer_type = "adam8bit"
 torch_compile = False
-unet_gradient = True
+unet_gradient = False
 clip_sample = False #Scheduler
 fixed_seed = False
 shuffle = True
@@ -400,7 +400,7 @@ if os.path.isdir(latest_checkpoint):
 if torch_compile:
     print("compiling")
     torch.set_float32_matmul_precision('high')
-    unet = torch.compile(unet)#, mode="reduce-overhead", fullgraph=True)
+    unet = torch.compile(unet, mode="reduce-overhead", fullgraph=False)
     print("compiling - ok")
 
 if lora_name: