add yet another function for automatically returning a finetune optimizer

Author: lucidrains

perfusion_pytorch/__init__.py

@@ -14,6 +14,10 @@
 
 from perfusion_pytorch.save_load import (
     save,
-    load,
-    get_finetune_parameters
+    load
+)
+
+from perfusion_pytorch.optimizer import (
+    get_finetune_parameters,
+    get_finetune_optimizer
 )

perfusion_pytorch/optimizer.py

@@ -0,0 +1,47 @@
+from torch.nn import Module
+from torch.optim import AdamW, Adam, Optimizer
+
+from beartype import beartype
+
+from perfusion_pytorch.embedding import EmbeddingWrapper
+from perfusion_pytorch.perfusion import Rank1EditModule
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+# function that automatically finds all the parameters necessary for fine tuning
+
+@beartype
+def get_finetune_parameters(text_image_model: Module):
+    params = []
+    for module in text_image_model.modules():
+        if isinstance(module, (EmbeddingWrapper, Rank1EditModule)):
+            params.extend(module.parameters())
+
+    return params
+
+@beartype
+def get_finetune_optimizer(
+    text_image_model: Module,
+    lr = 1e-4,
+    wd = 1e-2,
+    betas = (0.9, 0.99),
+    eps = 1e-8,
+    **kwargs
+) -> Optimizer:
+    params = get_finetune_parameters(text_image_model)
+
+    assert len(params) > 0, 'no finetuneable parameters found'
+    total_params = sum([p.numel() for p in params])
+    print(f'optimizer {total_params} parameters')
+
+    has_weight_decay = wd > 0
+    adam_klass = AdamW if has_weight_decay else Adam
+    adam_kwargs = dict(lr = lr, betas = betas, eps = eps)
+
+    if has_weight_decay:
+        adam_kwargs.update(weight_decay = wd)
+
+    return adam_klass(params, **adam_kwargs, **kwargs)

perfusion_pytorch/save_load.py

@@ -14,17 +14,6 @@
 def exists(val):
     return val is not None
 
-# function that automatically finds all the parameters necessary for fine tuning
-
-@beartype
-def get_finetune_parameters(text_image_model: Module):
-    params = []
-    for module in text_image_model.modules():
-        if isinstance(module, (EmbeddingWrapper, Rank1EditModule)):
-            params.extend(module.parameters())
-
-    return params
-
 # saving and loading the necessary extra finetuned params
 
 @beartype

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'perfusion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.17',
+  version = '0.1.18',
   license='MIT',
   description = 'Perfusion - Pytorch',
   author = 'Phil Wang',