sean-rice / srnet Goto Github PK

when using mnist_train to train classifier networks, we can see the following in the training log:

[12/29 17:13:19 d2.data.common]: Serializing 60000 elements to byte tensors and concatenating them all ...
[12/29 17:13:58 d2.data.common]: Serialized dataset takes 27546.78 MiB

of course, mnist should not be taking over 25GiB of memory. while mnist still works on the training cluster i use, this should be investigated and fixed; if the problem persists (and scales) to new datasets, it's possible other (still very small) datasets like cifar won't actually fit into worker memory.

The pre-processed *.pt files we use in srnet/data/datasets/mnist.py are no longer generated in newer versions of torchvision; see pytorch/vision#3420

This snippet might be equivalent:

# where `dataset_path` is `d2_datasets / "MNIST"` in the usual case, and `train` is a new parameter to `_make_mnist_dicts()`
mnist = MNIST(dataset_path, train=train)
data, labels_tensor = mnist.data, mnist.targets
data = data.to("cpu")
labels: Sequence[int] = labels_tensor.tolist()
del labels_tensor
n_examples, h, w, = data.shape
assert n_examples == len(labels)

pillow 10 removed some attributes we use.