TorchLearner: support / example for custom batch_sampler

[tdhock]

Related to #417

I expected to be able to use batch_sampler argument in mlr3torch, similar to how we can use it in torch.
However I observe that there are errors in mlr3torch, even using a batch_sampler that works with torch.

Here is an example of a batch_sampler that works in pure torch.

# 1. Define a simple dataset
my_dataset <- torch::dataset(
  name = "MyDataset",
  initialize = function() {
    self$data <- torch::torch_tensor(1:5)
  },
  .getitem = function(index) {
    self$data[index]
  },
  .length = function() {
    length(self$data)
  }
)
# 2. Define a custom batch sampler 
batch_sampler_class <- torch::sampler(
  "BatchSampler",
  initialize = function(data_source) {
    self$data_source <- data_source
  },
  .iter = function() {
    batch_size <- 2
    indices <- 1:self$.length()
    batch_vec <- (indices-1) %/% batch_size
    batch_list <- rev(split(indices, batch_vec))
    count <- 0L
    function() {
      if (count < length(batch_list)) {
        count <<- count + 1L
        return(batch_list[[count]])
      }
      coro::exhausted()
    }
  },
  .length = function() {
    length(self$data_source)
  }
)
# 3. Instantiate dataset and sampler
ds <- my_dataset()
batch_sampler_instance <- batch_sampler_class(ds)
# 4. Create dataloader with custom sampler
dl <- torch::dataloader(ds, batch_sampler = batch_sampler_instance)
# 5. Iterate through the dataloader
batches <- list()
coro::loop(for (batch in dl) {
  batches[[length(batches) + 1L]] = batch
})
batches

Running this code gives me the output below:

> batches
[[1]]
torch_tensor
 5
[ CPULongType{1} ]

[[2]]
torch_tensor
 3
 4
[ CPULongType{2} ]

[[3]]
torch_tensor
 1
 2
[ CPULongType{2} ]

The output above indicates that there are three batches. First with one sample, and the next two with two samples each.

Below I try something similar with mlr3torch current main.

sonar_task <- mlr3::tsk("sonar")
sonar_ingress <- list(feat=mlr3torch::TorchIngressToken(
  features=sonar_task$col_roles$feature,
  batchgetter=mlr3torch::batchgetter_num))
target_batchgetter <- function(data){
  browser()
  torch::torch_tensor()
}
sonar_dataset <- mlr3torch::task_dataset(sonar_task, sonar_ingress, target_batchgetter)
batch_sampler_instance <- batch_sampler_class(sonar_dataset)
inst_learner <- mlr3torch::LearnerTorchMLP$new(task_type="classif")
inst_learner$param_set$set_values(
  epochs=10,
  batch_size=20,
  batch_sampler=batch_sampler_instance)
inst_learner$train(sonar_task)

Note that in the code above I am using batch_sampler=batch_sampler_instance but it would be more convenient for benchmarks to provide batch_sampler=batch_sampler_class as proposed in #419.

I get the error below,

> inst_learner$train(sonar_task)
Error in ctx$batch$y$to(device = ctx$device) : 
  attempt to apply non-function
> traceback()
15: train_loop(ctx, callbacks)
14: learner_torch_train(self, private, super, task, param_vals)
13: force(expr)
12: with_torch_settings(seed = param_vals$seed, num_threads = param_vals$num_threads, 
        num_interop_threads = param_vals$num_interop_threads, expr = {
            learner_torch_train(self, private, super, task, param_vals)
        })
11: .__LearnerTorch__.train(self = self, private = private, super = super, 
        task = task)
10: get_private(learner)$.train(task)
9: .f(learner = <environment>, task = <environment>)
8: eval(expr, p)
7: eval(expr, p)
6: eval.parent(expr, n = 1L)
5: invoke(.f, .args = .args, .opts = .opts, .seed = .seed, .timeout = .timeout)
4: encapsulate(learner$encapsulation["train"], .f = train_wrapper, 
       .args = list(learner = learner, task = task), .pkgs = learner$packages, 
       .seed = NA_integer_, .timeout = learner$timeout["train"])
3: learner_train(learner, task, train_row_ids = train_row_ids, mode = mode)
2: .__Learner__train(self = self, private = private, super = super, 
       task = task, row_ids = row_ids)
1: inst_learner$train(sonar_task)

It seems that using a sampler instead of batch_sampler could be a work-around, but it would be inefficient (for loop over observations is much slower than for loop over batches).

@sebffischer Do you have any idea where this error is coming from? or how to fix?

[tdhock]

please re-open, this error is still happening with mlr3torch from CRAN (0.3.1) which includes #419

> inst_learner <- mlr3torch::LearnerTorchMLP$new(task_type="classif")
> inst_learner$param_set$set_values(
+  epochs=10,
+  batch_size=20,
+  batch_sampler=batch_sampler_class)
> inst_learner$train(sonar_task)
Error in ctx$batch$y$to(device = ctx$device) : 
  attempt to apply non-function

[sebffischer]

can you post a complete regex please

[tdhock]

here it is

## custom batch sampler
batch_sampler_class <- torch::sampler(
  "BatchSampler",
  initialize = function(data_source) {
    self$data_source <- data_source
  },
  .iter = function() {
    batch_size <- 2
    indices <- 1:self$.length()
    batch_vec <- (indices-1) %/% batch_size
    batch_list <- rev(split(indices, batch_vec))
    count <- 0L
    function() {
      if (count < length(batch_list)) {
        count <<- count + 1L
        return(batch_list[[count]])
      }
      coro::exhausted()
    }
  },
  .length = function() {
    length(self$data_source)
  }
)

## check that the batch_sampler works with torch
my_dataset <- torch::dataset(
  name = "MyDataset",
  initialize = function() {
    self$data <- torch::torch_tensor(1:5)
  },
  .getitem = function(index) {
    self$data[index]
  },
  .length = function() {
    length(self$data)
  }
)
ds <- my_dataset()
batch_sampler_instance <- batch_sampler_class(ds)
dl <- torch::dataloader(ds, batch_sampler = batch_sampler_instance)
batches <- list()
coro::loop(for (batch in dl) {
  batches[[length(batches) + 1L]] = batch
})
batches

## but it gives an error with mlr3torch
sonar_task <- mlr3::tsk("sonar")
inst_learner <- mlr3torch::LearnerTorchMLP$new(task_type="classif")
inst_learner$param_set$set_values(
  epochs=10,
  batch_size=20,
  batch_sampler=batch_sampler_class)
inst_learner$train(sonar_task)

[tdhock]

here is another positive control that shows that my custom batch sampler works with regular torch

## custom batch sampler
batch_sampler_class <- torch::sampler(
  "BatchSampler",
  initialize = function(data_source) {
    self$data_source <- data_source
  },
  .iter = function() {
    batch_size <- 2
    indices <- 1:self$.length()
    batch_vec <- (indices-1) %/% batch_size
    batch_list <- rev(split(indices, batch_vec))
    count <- 0L
    function() {
      if (count < length(batch_list)) {
        count <<- count + 1L
        return(batch_list[[count]])
      }
      coro::exhausted()
    }
  },
  .length = function() {
    length(self$data_source)
  }
)

## check that the batch_sampler works with torch and sonar data
sonar_task <- mlr3::tsk("sonar")
sonar_ingress <- list(feat=mlr3torch::TorchIngressToken(
  features=sonar_task$col_roles$feature,
  batchgetter=mlr3torch::batchgetter_num))
target_batchgetter <- function(data){
  torch::torch_tensor(ifelse(as.integer(data$Class)==1, 1, 0))
}
sonar_dataset <- mlr3torch::task_dataset(sonar_task, sonar_ingress, target_batchgetter)
batch_sampler_instance <- batch_sampler_class(sonar_dataset)
dl <- torch::dataloader(sonar_dataset, batch_sampler = batch_sampler_instance)
batches <- list()
coro::loop(for (batch in dl) {
  batches[[length(batches) + 1L]] = batch
})
str(batches[1:2])

I get the output below:

> str(batches[1:2])
List of 2
 $ :List of 3
  ..$ x     :List of 1
  .. ..$ feat:Float [1:2, 1:60]
  ..$ .index:Long [1:2]
  ..$ y     :Float [1:2]
 $ :List of 3
  ..$ x     :List of 1
  .. ..$ feat:Float [1:2, 1:60]
  ..$ .index:Long [1:2]
  ..$ y     :Float [1:2]

which indicates that the first two batches are each a list of 3 tensors.

[tdhock]

I instrumented mlr3torch:::train_loop with print/browser, 5da3414
And I observed that this error happens for the last batch, which is NULL. is that normal/expected? I guess not.

> inst_learner$train(sonar_task)
[1] 1
$x
$x$input
torch_tensor
Columns 1 to 10 0.0303  0.2354  0.2898  0.2812  0.1578  0.0273  0.0673  0.1444  0.2070  0.2645
 0.0260  0.2354  0.2720  0.2442  0.1665  0.0336  0.1302  0.1708  0.2177  0.3175

Columns 11 to 20 0.2828  0.0353  0.4293  0.5685  0.6990  0.7246  0.7622  0.9242  1.0000  0.9979
 0.3714  0.0363  0.4552  0.5700  0.7397  0.8062  0.8837  0.9432  1.0000  0.9375

Columns 21 to 30 0.8297  0.7032  0.0490  0.7141  0.6893  0.4961  0.2584  0.0969  0.0776  0.0364
 0.7603  0.7123  0.0136  0.8358  0.7622  0.4567  0.1715  0.1549  0.1641  0.1869

Columns 31 to 40 0.1572  0.1823  0.1349  0.0608  0.0849  0.0492  0.1367  0.1552  0.1548  0.1319
 0.2655  0.1713  0.0959  0.0272  0.0768  0.0847  0.2076  0.2505  0.1862  0.1439

Columns 41 to 50 0.0985  0.1258  0.0954  0.0489  0.0167  0.0241  0.0042  0.0086  0.0046  0.0126
 0.1470  0.0991  0.0041  0.0154  0.0214  0.0116  0.0181  0.0146  0.0129  0.0047

Columns 51 to 60 0.0036  0.0035  0.0034  0.0079  0.0036  0.1354  0.0048  0.1465  0.1123  0.1945
 0.0039  0.0061  0.0040  0.0036  0.0061  0.0338  0.0115  0.0655  0.1400  0.1843
[ CPUFloatType{2,60} ]


$.index
torch_tensor
 207
 208
[ CPULongType{2} ]

$y
torch_tensor
 1
 1
[ CPUFloatType{2,1} ]

...

[1] 104
$x
$x$input
torch_tensor
Columns 1 to 10 0.0200  0.2111  0.1609  0.1582  0.2238  0.0645  0.0660  0.2273  0.3100  0.2999
 0.0453  0.2872  0.4918  0.6552  0.6919  0.7797  0.7464  0.9444  1.0000  0.8874

Columns 11 to 20 0.5078  0.0371  0.4797  0.5783  0.5071  0.4328  0.5550  0.6711  0.6415  0.7104
 0.8024  0.0523  0.7818  0.5212  0.4052  0.3957  0.3914  0.3250  0.3200  0.3271

Columns 21 to 30 0.8080  0.6791  0.0428  0.3857  0.1307  0.2604  0.5121  0.7547  0.8537  0.8507
 0.2767  0.4423  0.0843  0.2028  0.3788  0.2947  0.1984  0.2341  0.1306  0.4182

Columns 31 to 40 0.6692  0.6097  0.4943  0.0207  0.2744  0.0510  0.2834  0.2825  0.4256  0.2641
 0.3835  0.1057  0.1840  0.0689  0.1970  0.1674  0.0583  0.1401  0.1628  0.0621

Columns 41 to 50 0.1386  0.1051  0.1343  0.0383  0.0954  0.0324  0.0232  0.0027  0.0065  0.0159
 0.0203  0.0530  0.0742  0.0409  0.1183  0.0061  0.0125  0.0084  0.0089  0.0048

Columns 51 to 60 0.0072  0.0167  0.0180  0.0084  0.0090  0.0986  0.0032  0.1539  0.1601  0.3109
 0.0094  0.0191  0.0140  0.0049  0.0052  0.2583  0.0044  0.2156  0.3481  0.3337
[ CPUFloatType{2,60} ]


$.index
torch_tensor
 1
 2
[ CPULongType{2} ]

$y
torch_tensor
 0
 0
[ CPUFloatType{2,1} ]

[1] 105
NULL
Called from: train_loop(ctx, callbacks)
Browse[1]> debug: ctx$batch$x = lapply(ctx$batch$x, function(x) x$to(device = ctx$device))
Browse[1]> c
Error: attempt to apply non-function
>

[tdhock]

the batch comes from dataloader_next which is documented on ?dataloader_next as

Usage:
     dataloader_next(iter, completed = NULL)
Arguments:
    iter: a DataLoader iter created with dataloader_make_iter.
completed: the returned value when the iterator is exhausted.

so returning NULL is expected behavior.
This suggests that mlr3torch code should be modified to handle that.

[tdhock]

the loop is while (ctx$step < length(ctx$loader_train)) { and these values in the failing iteration are

Browse[1]> ctx$step
[1] 105
Browse[1]> length(ctx$loader_train)
[1] 208

this suggests that the length is not being computed correctly. (should be 104 in this case)

[tdhock]

so my code .length() returns 208 which is the number of samples, consistent with documentation in ?torch::sampler

        * '.length' returns the maximum number of samples that can be
          retrieved from that sampler.

however it seems that .length() should return the number of batches (not samples), so maybe this is a bug in that documentation?
PyTorch docs say "Every Sampler subclass has to provide an __iter__() method, providing a way to iterate over indices or lists of indices (batches) of dataset elements, and may provide a __len__() method that returns the length of the returned iterators." https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.Sampler
which suggests that the R sampler docs are not correct.

[tdhock]

the fix is in my code: when .length() returns the number of batches (not samples), then the error goes away.

batch_sampler_class <- torch::sampler(
  "BatchSamplerWithLengthEqualToNumberOfBatches",
  initialize = function(data_source) {
    batch_size <- 2
    indices <- 1:length(data_source)
    batch_vec <- (indices-1) %/% batch_size
    self$batch_list <- rev(split(indices, batch_vec))
  },
  .iter = function() {
    count <- 0L
    function() {
      if (count < length(self$batch_list)) {
        count <<- count + 1L
        return(self$batch_list[[count]])
      }
      coro::exhausted()
    }
  },
  .length = function() {
    length(self$batch_list)
  }
)

The error message could be improved so I will file a PR.