library(luz)The Accelerator API is a simplified port of the Hugging Face Accelerate library. Currently it only handles CPU and single GPU usage but allows users avoid the boilerplate code necessary to write training loops that works correctly on both devices.
This API is meant to be the most flexible way you can use the luz package. With the Accelerator API, you write the raw torch training loop and with a few code changes you handle device placement of model, optimizers and dataloaders so you don’t need to add many $to(device="cuda") in your code or think about the order to create model and optimizers.
The Accelerator API is best explained by showing an example diff in a raw torch training loop.
library(torch)
+ library(luz)
+ acc <- accelerator()
- device <- "cpu"
data <- tensor_dataset(
x = torch_randn(100, 10),
y = torch_rand(100, 1)
)
dl <- dataloader(data, batch_size = 10)
model <- nn_linear(10, 1)
- model$to(device = device)
opt <- optim_adam(model$parameters)
+ c(model, opt, dl) %<-% acc$prepare(model, opt, dl)
model$train()
coro::loop(for (batch in dl) {
opt$zero_grad()
- preds <- model(batch$x$to(device = device))
+ preds <- model(batch$x)
- loss <- nnf_mse_loss(preds, batch$y$to(device = device))
+ loss <- nnf_mse_loss(preds, batch$y)
loss$backward()
opt$step()
})With the following changes to your code you no longer need to manually move data and parameters between devices which makes your code easier to read and less error prone.
You can find additional documentation using help(accelerator).