Skip to content

Fix grayscale channel dimension (1) lost in DataLoader #307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cregouby merged 6 commits into from
Mar 30, 2026
Merged

Fix grayscale channel dimension (1) lost in DataLoader #307

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e14d80b
Fix grayscale channel dimension (1) lost in DataLoader (#306)
Chandraveersingh1717 Mar 24, 2026
809c665
Merge branch 'main' into fix-grayscale-dimension
Chandraveersingh1717 Mar 25, 2026
1e50c94
fix requested changes
Chandraveersingh1717 Mar 28, 2026
c6e98bd
fix test failure
Chandraveersingh1717 Mar 28, 2026
44e4c87
align all .getitem() to be ndim = 3, .getbatch() to be ndim = 4 and d…
cregouby Mar 29, 2026
ec658a0
add NEWS
cregouby Mar 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 74 additions & 52 deletions R/dataset-mnist.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ mnist_dataset <- dataset(
archive <- download_and_cache(r[1], prefix = class(self)[1])
fs::file_copy(archive, destpath)

if (!tools::md5sum(destpath) == r[2])
if (!(tools::md5sum(destpath) == r[2]))
runtime_error("Corrupt file! Delete the file in {archive} and try again.")

}
Expand Down Expand Up @@ -149,8 +149,13 @@ mnist_dataset <- dataset(
},

.getitem = function(index) {
x <- self$data[index, ,]
y <- self$targets[index]
if (length(index) != 1) {
return(lapply(as.integer(index), function(i) self$.getitem(i)))
}

idx <- as.integer(index)

@cregouby cregouby Mar 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

todo Same as for L454-L458

x <- self$data[idx, , , drop = FALSE]
y <- self$targets[idx]

if (!is.null(self$transform))
x <- self$transform(x)
Expand Down Expand Up @@ -179,7 +184,7 @@ mnist_dataset <- dataset(
#' @describeIn mnist_dataset Kuzushiji-MNIST cursive Japanese character dataset.
#' @export
kmnist_dataset <- dataset(
name = "kminst_dataset",
name = "kmnist_dataset",
inherit = mnist_dataset,
archive_size = "21 MB",
resources = list(
Expand Down Expand Up @@ -252,55 +257,63 @@ qmnist_dataset <- dataset(
},

download = function() {
if (self$check_exists())
return(NULL)

fs::dir_create(self$raw_folder)
fs::dir_create(self$processed_folder)

cli_inform("Downloading {.cls {class(self)[[1]]}} ...")
for (r in self$resources[[self$split]]) {
filename <- basename(r[1])
destpath <- file.path(self$raw_folder, filename)

archive <- download_and_cache(r[1], prefix = glue::glue("qmnist-{self$split}"))
fs::file_copy(archive, destpath, overwrite = TRUE)

if (!tools::md5sum(destpath) == r[2])
runtime_error("Corrupt file! Delete the file in {archive} and try again.")
split_exists <- function(split_name) {
fs::file_exists(file.path(self$processed_folder, self$files[[split_name]]))
}

cli_inform("Processing {.cls {class(self)[[1]]}} ...")


if (self$split == "train") {
saveRDS(
list(
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-images-idx3-ubyte.gz")),
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-labels-idx2-int.gz"))
),
file.path(self$processed_folder, self$files$train)
)
}

if (self$split == "test") {
saveRDS(
list(
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-images-idx3-ubyte.gz")),
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-labels-idx2-int.gz"))
),
file.path(self$processed_folder, self$files$test)
)
}
if (all(vapply(names(self$resources), split_exists, logical(1))))
return(NULL)

if (self$split == "nist") {
saveRDS(
list(
read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-images-idx3-ubyte.xz")),
read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-labels-idx2-int.xz"))
),
file.path(self$processed_folder, self$files$nist)
)
cli_inform("Downloading {.cls {class(self)[[1]]}} ...")
for (split_name in names(self$resources)) {
if (split_exists(split_name))
next

for (r in self$resources[[split_name]]) {
filename <- basename(r[1])
destpath <- file.path(self$raw_folder, filename)

archive <- download_and_cache(r[1], prefix = glue::glue("qmnist-{split_name}"))
fs::file_copy(archive, destpath, overwrite = TRUE)

if (!(tools::md5sum(destpath) == r[2]))
runtime_error("Corrupt file! Delete the file in {archive} and try again.")
}

cli_inform("Processing {.val {split_name}} split for {.cls {class(self)[[1]]}} ...")

if (split_name == "train") {
saveRDS(
list(
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-images-idx3-ubyte.gz")),
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-labels-idx2-int.gz"))
),
file.path(self$processed_folder, self$files$train)
)
}

if (split_name == "test") {
saveRDS(
list(
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-images-idx3-ubyte.gz")),
read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-labels-idx2-int.gz"))
),
file.path(self$processed_folder, self$files$test)
)
}

if (split_name == "nist") {
saveRDS(
list(
read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-images-idx3-ubyte.xz")),
read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-labels-idx2-int.xz"))
),
file.path(self$processed_folder, self$files$nist)
)
}
}

cli_inform("Dataset {.cls {class(self)[[1]]}} downloaded and extracted successfully.")
Expand Down Expand Up @@ -388,7 +401,8 @@ emnist_collection <- dataset(
self$processed_folder <- file.path(root, class(self)[1], "processed")
self$transform <- transform
self$target_transform <- target_transform
self$class <- self$classes_all_dataset[[self$dataset]]
self$classes <- self$classes_all_dataset[[self$dataset]]
self$class <- self$classes

if (download) {
cli_inform("Dataset {.val {self$dataset}} split {.val {self$split}} of {.cls {class(self)[[1]]}} (~{.emph {self$archive_size}}) will be downloaded and processed if not already available.")
Expand All @@ -414,7 +428,7 @@ emnist_collection <- dataset(
url <- self$resources[[1]][1]
archive <- download_and_cache(url, prefix = class(self)[1])

if (!tools::md5sum(archive) == self$resources[[1]][2])
if (!(tools::md5sum(archive) == self$resources[[1]][2]))
runtime_error("Corrupt file! Delete the file in {archive} and try again.")

unzip_dir <- file.path(self$raw_folder, "unzipped")
Expand All @@ -437,9 +451,13 @@ emnist_collection <- dataset(
},

.getitem = function(index) {
if (length(index) != 1) {
return(lapply(as.integer(index), function(i) self$.getitem(i)))
}

@cregouby cregouby Mar 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion this is supposed to be the .getbatch() function. I would rather implement the .getbatch than trying to implement a get-batch-within-getitem.


x <- self$data[index, , ]
y <- self$targets[index]
idx <- as.integer(index)

@cregouby cregouby Mar 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion `.getitem() is mainly an internal function for dataloader to run. So index will always be integer. I'd prefer to avoid the perrformance impact and remove this line.

x <- self$data[idx, , , drop = FALSE]
y <- self$targets[idx]

if (!is.null(self$transform))
x <- self$transform(x)
Expand All @@ -457,7 +475,11 @@ emnist_collection <- dataset(
)

read_sn3_pascalvincent <- function(path) {
x <- gzfile(path, open = "rb")
x <- if (grepl("\\\\.xz$", path, ignore.case = TRUE)) {
xzfile(path, open = "rb")
} else {
gzfile(path, open = "rb")
}
Comment on lines +477 to +481

@cregouby cregouby Mar 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

praise nice addition

on.exit({close(x)})

magic <- readBin(x, endian = "big", what = integer(), n = 1)
Expand Down
11 changes: 10 additions & 1 deletion R/transforms-array.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,16 @@ transform_to_tensor.array <- function(img) {
if (length(dim(img)) == 2)
dim(img) <- c(dim(img), 1)

res <- torch::torch_tensor(img)$permute(c(3, 1, 2))
dims <- dim(img)
if (length(dims) != 3)
stop("Expected a 2D or 3D array.")

@cregouby cregouby Mar 28, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

todo Please rely on functions from conditions.R for correct user message translation


# Support both HWC (default image arrays) and CHW (channel-first arrays).
if (dims[1] <= 4 && dims[3] > 4) {
res <- torch::torch_tensor(img)
} else {
res <- torch::torch_tensor(img)$permute(c(3, 1, 2))
}

if (res$dtype == torch::torch_long())
res <- res/255
Expand Down
26 changes: 21 additions & 5 deletions tests/testthat/test-dataset-mnist.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,21 @@ test_that("tests for the mnist dataset", {
ds <- mnist_dataset(dir, download = TRUE)

i <- ds[1]
expect_equal(dim(i[[1]]), c(28, 28))
expect_equal(dim(i[[1]]), c(1, 28, 28))
expect_equal(i[[2]], 6)
expect_length(ds, 60000)

raw_items <- ds$.getitem(c(1, 2))
expect_length(raw_items, 2)
expect_named(raw_items[[1]], c("x", "y"))
expect_equal(dim(raw_items[[1]]$x), c(1, 28, 28))

ds <- mnist_dataset(dir, transform = transform_to_tensor)
items <- ds$.getitem(c(1, 2))
expect_length(items, 2)
expect_named(items[[1]], c("x", "y"))
expect_tensor_shape(items[[1]]$x, c(1, 28, 28))

dl <- torch::dataloader(ds, batch_size = 32)
expect_length(dl, 1875)
iter <- dataloader_make_iter(dl)
Expand All @@ -25,6 +35,12 @@ test_that("tests for the mnist dataset", {
expect_true((torch_max(i[[1]]) <= 1)$item())
expect_named(i, c("x", "y"))

# Regression: grayscale channel must be preserved for larger batches.
dl_128 <- torch::dataloader(ds, batch_size = 128)
i_128 <- dataloader_next(dataloader_make_iter(dl_128))
expect_tensor_shape(i_128$x, c(128, 1, 28, 28))
expect_tensor_shape(i_128$y, 128)

})

test_that("tests for the kmnist dataset", {
Expand All @@ -36,7 +52,7 @@ test_that("tests for the kmnist dataset", {
ds <- kmnist_dataset(dir, download = TRUE)

i <- ds[1]
expect_equal(dim(i[[1]]), c(28, 28))
expect_equal(dim(i[[1]]), c(1, 28, 28))
expect_equal(i[[2]], 6)
expect_length(ds, 60000)

Expand All @@ -62,7 +78,7 @@ test_that("fashion_mnist_dataset loads correctly", {
expect_s3_class(ds, "fashion_mnist_dataset")
expect_type(ds$.getitem(1), "list")
expect_named(ds$.getitem(1), c("x", "y"))
expect_equal(dim(as.array(ds$.getitem(1)$x)), c(28, 28))
expect_equal(dim(as.array(ds$.getitem(1)$x)), c(1, 28, 28))
expect_true(ds$.getitem(1)$y >= 1 && ds$.getitem(1)$y <= 10)

ds2 <- fashion_mnist_dataset(dir, transform = transform_to_tensor)
Expand Down Expand Up @@ -96,7 +112,7 @@ test_that("tests for the emnist dataset", {
first_item <- emnist[1]
expect_named(first_item, c("x", "y"))
expect_s3_class(first_item$x, "array")
expect_equal(dim(first_item$x), c(28,28))
expect_equal(dim(first_item$x), c(1, 28, 28))
expect_equal((first_item[[2]]), 19)

emnist <- emnist_collection(dir, dataset = "bymerge", download = TRUE)
Expand Down Expand Up @@ -157,7 +173,7 @@ test_that("tests for the qmnist dataset", {
ds <- qmnist_dataset(dir, split = split, download = TRUE)

i <- ds[1]
expect_equal(dim(i[[1]]), c(28, 28))
expect_equal(dim(i[[1]]), c(1, 28, 28))
expect_true(i[[2]] %in% 1:10)

expect_gt(length(ds), 0)
Expand Down
Loading