michael/modrec-workflow
1
0
Fork 0
forked from qoherent/modrec-workflow
Code Pull Requests Workflows Activity

changed back to gpu

Browse Source
This commit is contained in:
Liyu Xiao 2025-05-26 14:38:18 -04:00
parent 92a0ed11e4
commit d518c633b7
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
data/training/train.py
View File

@ -118,12 +118,13 @@ def train_model():
drop_path_rate=hparams["drop_path_rate"], drop_path_rate=hparams["drop_path_rate"],
) )
) )
checkpoint_callback = ModelCheckpoint( checkpoint_callback = ModelCheckpoint(
dirpath=checkpoint_dir, dirpath=checkpoint_dir,
filename=checkpoint_filename, filename=checkpoint_filename,
save_top_k=1, save_top_k=1,
verbose=False, # Disable checkpoint verbose output verbose=True, # Disable checkpoint verbose output
monitor="val_acc", monitor="val_acc",
mode="max", mode="max",
enable_version_counter=False, enable_version_counter=False,
@ -132,10 +133,10 @@ def train_model():
trainer = L.Trainer( trainer = L.Trainer(
max_epochs=epochs, max_epochs=epochs,
callbacks=[checkpoint_callback], callbacks=[checkpoint_callback],
accelerator="cpu", accelerator="gpu",
devices=1, devices=1,
benchmark=True, benchmark=True,
precision="bf16-mixed", precision="16-mixed",
logger=False, logger=False,
enable_progress_bar=False, # Disable all progress bars enable_progress_bar=False, # Disable all progress bars
enable_model_summary=False, # Disable model summary enable_model_summary=False, # Disable model summary