ComfyUI Node

FlattenDataset

Flatten your whole dataset's pixels in one move (with a catch)

By TashaSkyUp·Created about a year ago·Updated about a year ago· 1
FlattenDataset
  • dataset
  • TORCH_TENSOR

When you download MNIST with this pack's PyTorchDatasetDownloader, every sample is a 28×28 image. A dense network wants 784 flat numbers. FlattenDataset turns the whole dataset's pixel data into one big flattened tensor in a single step - torch.flatten(dataset.data) - so you can inspect it, feed it to a model, or build tensors from it without flattening sample-by-sample.

It's a convenience node for the download-and-inspect phase of a workflow. You grab a dataset, flatten its data, and use the result where you'd use any TORCH_TENSOR - say, checking shapes, or feeding a simple model that expects flat input.

The catch

The node does dataset.data. That attribute exists on raw torchvision datasets like MNIST and CIFAR (where .data is the full tensor of images). It does not exist on the pack's DatsetSplitter output - that returns torch.utils.data.Subset objects, which have no .data attribute and will raise an AttributeError here. So the sensible order is: flatten the dataset before splitting, or flatten the raw download. This is the classic "it works in the tutorial and breaks in my workflow" node.

Input and output

  • dataset (TORCH_DATASET) - a raw dataset object with a .data attribute, e.g. the output of PyTorchDatasetDownloader in dataset mode.
  • Output: one TORCH_TENSOR - the entire dataset's data flattened into a single 1D tensor.

Install

ComfyUI Manager, search "EternalKernel PyTorch Nodes", or:

cd ComfyUI/custom_nodes
git clone https://github.com/TashaSkyUp/EternalKernelPytorchNodes
cd EternalKernelPytorchNodes
pip install -r requirements.txt

Restart ComfyUI; node under ETK/pytorch. No model files. Requirements are the standard stack plus scipy, scikit-learn, transformers, einops.

Common issues

  • AttributeError: 'Subset' object has no attribute 'data'. The #1 failure, caused by feeding a split dataset. Flatten before splitting.
  • It flattens the labels too, if present. dataset.data on torchvision datasets is just the inputs - but be aware you're flattening the data field, not the dataset as a whole. If you need labels alongside, grab the labels tensor output of the downloader separately.
  • Giant tensors. Flattening 60,000 MNIST images gives you a ~47-million-element tensor. Fine in RAM, but if you then feed the whole thing to a batch-size-1 TrainModel it'll be slow - split and batch first.

Honest assessment: for most workflows you'll reach for FlattenTensor (per-tensor, cleaner) rather than this node. FlattenDataset is for the quick-and-dirty path where you want the dataset's pixels as one flat blob. The pack has no community tutorials, and the pack-wide quirk (it patches ComfyUI's validator to ignore return_type_mismatch) means you should verify wire types yourself when outputs look wrong.

CategoryETK/pytorch

Inputs (1)

NameTypeDefaultDescription
datasetTORCH_DATASET

Outputs (1)

NameTypeDescription
TORCH_TENSORTORCH_TENSOR