diff --git a/pytorch_nodes.py b/pytorch_nodes.py index e387372..666e455 100644 --- a/pytorch_nodes.py +++ b/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,) + diff --git a/tests/test_pytorch_nodes.py b/tests/test_pytorch_nodes.py index fca1e43..51815df 100644 --- a/tests/test_pytorch_nodes.py +++ b/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())