Add batchnorm, dropout, and trainable nodes

pytorch_nodes.py

@@ -1296,3 +1296,84 @@ class PlotSeriesString:
         buf.close()
         return (im,)
 
+
+@ETK_pytorch_base
+class AddBatchNormLayer:
+    """Adds a BatchNorm layer to an ``nn.Sequential`` model."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("TORCH_MODEL",),
+                "num_features": ("INT", {"min": 1}),
+                "one_d": ("BOOLEAN", {"default": True}),
+            }
+        }
+
+    RETURN_TYPES = ("TORCH_MODEL",)
+    FUNCTION = "add_batchnorm_layer"
+    CATEGORY = "model"
+
+    def add_batchnorm_layer(self, model, num_features, one_d=True):
+        layer = nn.BatchNorm1d(num_features) if one_d else nn.BatchNorm2d(num_features)
+        if isinstance(model, nn.Sequential):
+            model.insert(len(model), layer)
+        else:
+            raise TypeError("The provided model is not a nn.Sequential model.")
+        return (model,)
+
+
+@ETK_pytorch_base
+class AddDropoutLayer:
+    """Adds a Dropout layer to an ``nn.Sequential`` model."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("TORCH_MODEL",),
+                "p": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.5}),
+            },
+            "optional": {
+                "inplace": (["True", "False"], {"default": "False"}),
+            },
+        }
+
+    RETURN_TYPES = ("TORCH_MODEL",)
+    FUNCTION = "add_dropout_layer"
+    CATEGORY = "model"
+
+    def add_dropout_layer(self, model, p=0.5, inplace="False"):
+        inplace = True if inplace in (True, "True") else False
+        layer = nn.Dropout(p=p, inplace=inplace)
+        if isinstance(model, nn.Sequential):
+            model.insert(len(model), layer)
+        else:
+            raise TypeError("The provided model is not a nn.Sequential model.")
+        return (model,)
+
+
+@ETK_pytorch_base
+class SetModelTrainable:
+    """Sets ``requires_grad`` for all parameters of a model."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("TORCH_MODEL",),
+                "requires_grad": (["True", "False"], {"default": "False"}),
+            }
+        }
+
+    RETURN_TYPES = ("TORCH_MODEL",)
+    FUNCTION = "set_trainable"
+    CATEGORY = "model"
+
+    def set_trainable(self, model, requires_grad="False"):
+        flag = True if requires_grad in (True, "True") else False
+        for param in model.parameters():
+            param.requires_grad = flag
+        return (model,)
+

tests/test_pytorch_nodes.py

@@ -36,6 +36,9 @@ from custom_nodes.EternalKernelLiteGraphNodes.pytorch_nodes import (
     FuncModifyModel,
     PlotSeriesString,
     GridSearchTraining,
+    AddBatchNormLayer,
+    AddDropoutLayer,
+    SetModelTrainable,
 )
 
 
@@ -318,3 +321,26 @@ def test_grid_search_training(monkeypatch):
     best_model, best_params, metrics = node.grid_search_train(model, param_grid)
     assert best_params["epochs"] == 1
     assert len(metrics) == 2
+
+
+def test_add_batchnorm_layer():
+    model = nn.Sequential(nn.Linear(4, 4))
+    node = AddBatchNormLayer()
+    out_model, = node.add_batchnorm_layer(model, 4, True)
+    assert isinstance(out_model[-1], nn.BatchNorm1d)
+
+
+def test_add_dropout_layer():
+    model = nn.Sequential(nn.Linear(4, 4))
+    node = AddDropoutLayer()
+    out_model, = node.add_dropout_layer(model, 0.3)
+    assert isinstance(out_model[-1], nn.Dropout)
+
+
+def test_set_model_trainable():
+    model = nn.Sequential(nn.Linear(1, 1))
+    node = SetModelTrainable()
+    frozen, = node.set_trainable(model, "False")
+    assert all(not p.requires_grad for p in frozen.parameters())
+    unfrozen, = node.set_trainable(frozen, "True")
+    assert all(p.requires_grad for p in unfrozen.parameters())