Model_FireType_M3.R

library(torch)
library(torchvision)



# ----- Torch utilities -----



# Dataset class
ndvi_dataset <- torch::dataset(
  
  # Dataset name
  name = "ndvi_dataset",
  
  # Initialization function
  # Extracts the image paths to pass to the loader
  # Matches the images with their target info if it is available
  initialize = function(
    target_path,
    img_dir
  ) {
    
    # Load the image paths
    img_paths <- list.files(
      img_dir,
      full.names = TRUE,
      recursive = FALSE
    )
    img_id <- gsub("(^.+?_)(r[0-9]{5}.*)", "\\2", basename(img_paths))
    
    # Load the target data
    target_data <- data.table::fread(
      target_path,
      data.table = FALSE
    )
    target_id <- sprintf("r%05d_c%05d.png", target_data$row_id, target_data$col_id)
    
    # Format the target info for the loaded images
    y <- target_data[, grepl("^[A-Z]{2}", names(target_data))]
    y <- as.matrix(y[match(img_id, target_id),])
    
    # Add all observations with non-zero target info
    all_zero <- (rowSums(y) == 0)
    warning(sum(all_zero), " observations have no target info and were removed.")
    self$images <- data.frame(
      img = img_paths[!all_zero],
      img_id = img_id[!all_zero]
    )
    self$y <- torch::torch_tensor(y[!all_zero,], dtype = torch::torch_float()) 
    
  },
  
  # Get an item from the dataset
  # Only keep one of the channels since they are all identical
  .getbatch = function(index){
    img_batch <- lapply(index, function(i){
      img <- torchvision::base_loader(self$images$img[i])
      torchvision::transform_to_tensor(img[, , 1])
    })
    img_batch <- torch::torch_stack(img_batch, 1)
    y_batch <- self$y[index,]
    return(list(x = img_batch, y = y_batch))
  },
  
  # Return the dataset size
  .length = function() {
    nrow(self$images)
  }
  
)



# ----- Model -----



# Convolutional block
conv_block <- torch::nn_module( 
  "conv_block",
  initialize = function(in_size, out_size) {
    self$conv_block <- torch::nn_sequential(
      torch::nn_conv2d(in_size, out_size, kernel_size = 3, padding = "same"),
      torch::nn_relu()
    )
  },
  forward = function(x){
    self$conv_block(x)
  }
)

# Residual block
resid_block <- torch::nn_module( 
  "resid_block",
  initialize = function(in_size, out_size) {
    self$conv_1 <- torch::nn_sequential(
      torch::nn_conv2d(in_size, out_size, kernel_size = 3, padding = "same"),
      torch::nn_batch_norm2d(out_size),
      torch::nn_relu()
    )
    self$conv_2 <- torch::nn_sequential(
      torch::nn_conv2d(out_size, out_size, kernel_size = 3, padding = "same"),
      torch::nn_batch_norm2d(out_size)
    )
    self$relu <- torch::nn_relu()
  },
  forward = function(x){
    x_tilde <- x
    out <- self$conv_1(x)
    out <- self$conv_2(out)
    self$relu(x_tilde + out)
  }
)

# Main CNN
ndvi_cnn <- torch::nn_module(
  
  # Model name
  "ndvi_cnn",
  
  # Initialization
  initialize = function(
    channels_in,
    n_classes = 7
  ) {
    
    # Simple Resnet
    self$cnn <- torch::nn_module_list()
    self$cnn$append(resid_block(channels_in, 60))
    self$cnn$append(torch::nn_avg_pool2d(kernel_size = 4))
    self$cnn$append(torch::nn_conv2d(60, 30, 1))
    self$cnn$append(resid_block(30, 30))
    self$cnn$append(torch::nn_conv2d(30, 10, 1))
    self$cnn$append(torch::nn_avg_pool2d(kernel_size = 4))
    
    # MLP
    self$mlp <- torch::nn_module_list()
    self$mlp$append(torch::nn_flatten(start_dim = 2, end_dim = -1))
    self$mlp$append(torch::nn_linear(in_features = 4 * 4 * 10, out_features = 49))
    self$mlp$append(torch::nn_dropout(0.6))
    self$mlp$append(torch::nn_batch_norm1d(49))
    self$mlp$append(torch::nn_relu())
    self$mlp$append(torch::nn_linear(in_features = 49, out_features = n_classes))
    self$mlp$append(torch::nn_softmax(dim = 2))
    
  },
  
  forward = function(x) {
    for(i in 1:length(self$cnn)){
      #message("CNN module ", i)
      #message(paste0(dim(x), collapse = " "))
      x <- self$cnn[[i]](x)
    }
    #message("CNN out:", paste0(dim(x), collapse = " "))
    for(i in 1:length(self$mlp)){
      #message("MLP module ", i)
      #message(paste0(dim(x), collapse = " "))
      x <- self$mlp[[i]](x)
    }
    return(x)
  }
  
)

# Ramanujan's approximation for log(x!)
rj_log_factorial <- function(x){
  num_1 <- (x + 4 * torch::torch_pow(x, 2) + 8 * torch::torch_pow(x, 3))
  term_1 <- (1 / 6) * torch::torch_log(num_1)
  term_2 <- (log(pi) / 2)
  return(x * torch::torch_log(x) - x + term_1 + term_2)
}

# Loss function
# Multinomial log loss (negative)
# Calculates as though one of each class has been observed 
# to avoid the undefined case for rj_log_factorial(0)
multi_ll <- function(output, target){
  n <- torch::torch_sum(target, dim = 2) + dim(target)[2]
  log_n_fac <- rj_log_factorial(n)
  sum_log_y_fac <- torch::torch_sum(rj_log_factorial(target + 1), dim = 2)
  sum_ylog_p <- torch::torch_sum((target + 1) * torch::torch_log(output), dim = 2)
  ll <- log_n_fac - sum_log_y_fac + sum_ylog_p
  return(torch::torch_sum(ll))
}



# ----- Testing -----



# Specify the computational device
device <- torch::torch_device("cuda")

# Initialize the data and assign a dataloader
ndvi_data <- ndvi_dataset(
  target_path = "ProcessedData/20250421_fuelClass.csv", 
  img_dir = "C:/Users/CBuglino/Desktop/NDVI_Data"
)
batch_size <- 2^7
data_train <- torch::dataloader(ndvi_data, batch_size, shuffle = TRUE)

# Initialize the model, optimizer, and learning rate scheduler
model <- ndvi_cnn(channels_in = 1, n_classes = 7)
model$to(device = device, copy = TRUE)
optimizer <- torch::optim_adam(model$parameters, lr = 0.001)
scheduler <- torch::lr_step(optimizer, step_size = 30, gamma = 0.1)

# Precision for the Gaussian weight prior
w_var <- (2)^2
w_alpha <- 1 / w_var


# Train
epochs <- 125
epoch_ll <- rep(NA, epochs)
for(i in 1:epochs){
  
  # Mini batch learning
  message(format(Sys.time(), "%H:%M:%S"), " - Epoch: ", i)
  ll_list <- list()
  model$train()
  coro::loop(for(batch in data_train){
    
    # Compute maximum likelihood
    x <- batch$x$to(device = device)
    p_hat <- model(x)
    y <- batch$y
    loss <- -1 * multi_ll(p_hat$to(device = "cpu"), y)
    
    # Add Gaussian prior to the weights
    w_sq <- sapply(model$parameters, function(x){
      torch::torch_sum(torch::torch_pow(x, 2))
    })
    wtw <- torch::torch_sum(torch::torch_stack(w_sq))$to(device = "cpu")
    loss <- loss + (w_alpha / 2) * wtw
    
    if(as.numeric(torch::torch_isnan(loss)) == 1){
      stop()
    }
    
    # Update step
    optimizer$zero_grad()
    loss$backward()
    optimizer$step()
    ll_list[[length(ll_list) + 1]] <- as.numeric(-loss$detach())
    
  })
  model$eval()
  
  # Record total log likelihood
  ll <- sum(unlist(ll_list))
  epoch_ll[i] <- ll
  
  # Check Training
  if(i %% 1 == 0){
    ll <- sum(unlist(ll_list))
    cat(" - Log likelihood:", ll, "\n")
  }
  if((i %% 5 == 0) | (i == 1)){
    nom <- sprintf("model_e%03d.pt", i)
    torch::torch_save(model, path = nom)
  }
  
  # Update learning rate
  scheduler$step()
  
}