Add .getbatch to dataloaders that inherits mnist_dataset

NEWS.md

@@ -1,5 +1,7 @@
 # luz (development version)
 
+* override `.getbatch()` method in examples with `mnist_dataset`. (@cregouby #160)
+
 # luz 0.5.1
 
 * fixed a bug preventing additional arguments passed to `predict` to be forwarded to `model$predict` (#157)

vignettes/examples/mnist-autoencoder.Rmd

@@ -21,8 +21,12 @@ mnist_dataset2 <- torch::dataset(
     output <- super$.getitem(i)
     output$y <- output$x
     output
-  }
-)
+  },
+  .getbatch = function(i) {
+    output <- super$.getbatch(i)
+    output$y <- output$x
+    output
+  })
 
 train_ds <- mnist_dataset2(
   dir,

vignettes/examples/mnist-triplet.Rmd

@@ -16,6 +16,11 @@ dir <- "./mnist" # caching directory
 
 triplet_mnist_dataset <- dataset(
   inherit = mnist_dataset,
+  initialize = function(...) {
+    super$initialize(...)
+    self$targets_vec <- as.integer(self$targets)
+  },
+
   .getitem = function(index) {
     anchor <- self$data[index, ,]
     label <- self$targets[index]
@@ -31,6 +36,25 @@ triplet_mnist_dataset <- dataset(
       ),
       list() # no 'target'
     )
+  },
+
+  .getbatch = function(indices) {
+    safe_sample <- function(x) {
+      if (length(x) <= 1) return(x)
+      sample(x, 1)
+    }
+    labels <- self$targets_vec[indices]
+    pos_idx <- vapply(labels, function(l) safe_sample(which(self$targets_vec == l)), 1L)
+    neg_idx <- vapply(labels, function(l) safe_sample(which(self$targets_vec != l)), 1L)
+
+    list(
+      list( # input is a list with 3 images.
+        anchor = self$transform(self$data[indices, , ])$permute(c(2,1,3))$unsqueeze(2),
+        negative = self$transform(self$data[neg_idx, , ])$permute(c(2,1,3))$unsqueeze(2),
+        positive = self$transform(self$data[pos_idx, , ])$permute(c(2,1,3))$unsqueeze(2)
+      ),
+      list() # no 'target'
+    )
   }
 )