ComfyUI Node
TrainModel
A ComfyUI node in ETK/pytorch with 11 inputs and 3 outputs.
TrainModel
- model
- dataset
- features tensor
- labels tensor
- model
- metrics
- best model
◄epochs1►
◄batch_size1►
◄loss_functionCrossEntropyLoss►
◄optimizerSGD►
◄metricsaccuracy►
◄devicetorch.device('cpu')►
◄create_samplesFALSE►
CategoryETK/pytorch
Inputs (11)
| Name | Type | Default | Description |
|---|---|---|---|
| model | TORCH_MODEL | — | |
| datasetopt | TORCH_DATASET | — | |
| features tensoropt | TORCH_TENSOR | — | |
| labels tensoropt | TORCH_TENSOR | — | |
| epochsopt | INT | 11–16777216 | — |
| batch_sizeopt | INT | 1 | — |
| loss_functionopt | COMBO | CrossEntropyLoss | 22 options: AdaptiveLogSoftmaxWithLoss, BCELoss, BCEWithLogitsLoss, CTCLoss, CosineEmbeddingLoss, CrossEntropyLoss, +16 |
| optimizeropt | COMBO | SGD | 15 options: ASGD, Adadelta, Adafactor, Adagrad, Adam, AdamW, +9 |
| metricsopt | COMBO | accuracy | 23 options: accuracy, AdaptiveLogSoftmaxWithLoss, BCELoss, BCEWithLogitsLoss, CTCLoss, CosineEmbeddingLoss, +17 |
| deviceopt | STRING | torch.device('cpu') | — |
| create_samplesopt | COMBO | FALSE | 2 options: TRUE, FALSE |
Outputs (3)
| Name | Type | Description |
|---|---|---|
| model | TORCH_MODEL | — |
| metrics | LIST | — |
| best model | TORCH_MODEL | — |