---
title: 'Homework #5: Using a pretrained neural network'
output:
  html_document:
    df_print: paged
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
set.seed(24)

# Helper function for plotting images
plot_images <- function(batch) {
  images <- as_array(batch[[1]]) %>% aperm(perm = c(1, 3, 4, 2))
  images <- images * 255
  images[images > 255] <- 255
  images[images < 0] <- 0
  
  par(mfcol = c(4,6), mar = rep(1, 4))
  
  images %>%
    purrr::array_tree(1) %>%
    purrr::set_names(as_array(batch[[2]])) %>%
    purrr::map(as.raster, max = 255) %>%
    purrr::iwalk(~{plot(.x); title(.y)})  
}

```

1. Install and load the `torch` and `torchvision` packages.
```{r installTorch}
library(torch)
library(torchvision)
```

2. (1 point) Use the `tiny_imagenet_dataset` method to download a tiny version of the ImageNet dataset and name the dataset `train_ds`. Modify the template code as needed. Explain in your own words what the `transform` argument does. (When you finally knit this Rmd file, you may need to set `download=FALSE` rather than `download=TRUE`)
```{r imagenet}
train_ds <- tiny_imagenet_dataset(
  "~/Downloads",
  # set this to download = TRUE the first time you run this code
  download = FALSE,
  split="train",
  transform = function(x) {
    x %>%
      transform_to_tensor() %>%
      transform_resize(c(224, 224))
  }
)
```

3. (1 point) Create a vector `train_subset_ind` containing the index of all observations in `train_ds` with class labels 1, 2, or 3. The class labels are stored in the `targets` attribute.
```{r}
NUM_CLASSES <- 3
train_subset_ind <- as.numeric(which(train_ds$targets <= NUM_CLASSES))
```

4. (1 point) Create a subset of the original dataset composed of only observations with labels 1, 2, or 3 using the `dataset_subset` method and the vector `train_subset_ind` from question 3.
```{r}
train_small <- dataset_subset(train_ds, indices=train_subset_ind)
```

5. (2 points) Randomly select 20 images from the dataset and visualize them. To do this, create a data loader with batch size 20 with `shuffle=TRUE` and name it `train_dl_shuffle`. Read the first batch of data from `train_dl_shuffle` and use the `plot_images` function to plot the images.
```{r}
train_dl_temp <- dataloader(train_small, batch_size=20, shuffle = T)
batch <- train_dl_temp$.iter()$.next()

plot_images(batch)
```

6. (2 points) Create a new data loader for `train_small` that does not shuffle the data. Set the batch size to the maximum batch size. Name this new data loader `train_dl`.
```{r}
ntrain <- length(train_subset_ind)
train_dl <- dataloader(train_small, batch_size=ntrain, shuffle = F)
```

7. (3 points) Load a pretrained ResNet-18 model using the function `model_resnet18` with `pretrained=TRUE`. Name the pretrained model `pretrained_resnet`. List the types of layers used in ResNet-18.
```{r}
pretrained_resnet <- model_resnet18(pretrained = TRUE)
# This lists the major modules in Resnet-18. Note that `layer1` to `layer4` contain many layers within them.
pretrained_resnet
# To see what layers are used in `layer1`, run the following:
pretrained_resnet$layer1$`0`
pretrained_resnet$layer1$`1`
```

## Extracting features using a pretrained neural network

8. (1 point) Our goal is to train a new model using the extracted features by ResNet-18. We'll use the penultimate layer as the extracted features. To get extracted features, replace the last layer in `pretrained_resnet` with an identity module. That is, set the `fc` attribute in `pretrained_resnet` to a newly created identity module using the `nn_identity` function.
```{r}
pretrained_resnet$fc <- nn_identity()
```

9. (1 point) Set the pretrained ResNet-18 model to evaluation mode by calling the `eval` method of the `pretrained_resnet` object.
```{r}
pretrained_resnet$eval()
```

10. (1 point) Read the first batch of data from the data loader.
```{r}
my_data_obs <- train_dl$.iter()$.next()
```

11. (2 points) Feed this batch of data through our (modified) pretrained network `pretrained_resnet` to get the values of the penultimate layer. What is the dimension of this output? Print the embedding for the first observation.
```{r}
penultimate_dat_obs <- pretrained_resnet(my_data_obs[[1]])
# Dimension of the embedding layer
penultimate_dat_obs$size()[2]
# Embedding for the first observation
as_array(penultimate_dat_obs[1,])
```

## Training a model using extracted features

12. (2 points) Create a data frame with the embeddings from ResNet-18 and their corresponding class labels.
```{r}
my_df <- data.frame(as_array(penultimate_dat_obs), y = as.numeric(train_ds$targets[train_subset_ind]))
```

13. (2 points) Split this data so that 4/5 of the data is used for training and 1/5 of it is used for testing.
```{r}
my_train_idxs <- sample(nrow(my_df), size=nrow(my_df) * 4/5, replace=F)
my_df_train <- my_df[my_train_idxs,]
my_df_test <- my_df[-my_train_idxs,]
```

14. (2 points) Fit a random forest for this multiclassification problem using the training data. Use the default hyperparameter settings.
```{r}
library(randomForest)
rf <- randomForest(as.factor(y) ~ ., data=my_df_train)
rf
```

15. (1 point) Evaluate the random forest model on the test data. What is the accuracy of this model?
```{r}
preds <- predict(rf, newdata = my_df_test)
mean(preds == as.factor(my_df_test$y))
```