Add LoRA support for livecell finetuning scripts (#681)

computational-cell-analytics · Sep 5, 2024 · ff0539b · ff0539b
1 parent d0a31eb
commit ff0539b
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 139 deletions.
diff --git a/finetuning/livecell/lora/README.md b/finetuning/livecell/lora/README.md
diff --git a/finetuning/livecell/lora/train_livecell.py b/finetuning/livecell/lora/train_livecell.py
diff --git a/finetuning/livecell_finetuning.py b/finetuning/livecell_finetuning.py
@@ -72,6 +72,7 @@ def finetune_livecell(args):
         save_root=args.save_root,
         scheduler_kwargs=scheduler_kwargs,
         save_every_kth_epoch=args.save_every_kth_epoch,
+	lora_rank=args.lora_rank,
     )
 
     if args.export_path is not None:
@@ -116,9 +117,13 @@ def main():
     parser.add_argument(
         "--n_objects", type=int, default=25, help="The number of instances (objects) per batch used for finetuning."
     )
+    parser.add_argument(
+    	"--lora_rank", type=int, default=None, help="The rank for low rank adaptation of the attention layers."
+    )
     args = parser.parse_args()
     finetune_livecell(args)
 
 
 if __name__ == "__main__":
     main()
+