Forem: Favil Orbedios

fast.ai Book in Rust - Chapter 2 - Part 2

Favil Orbedios — Fri, 01 Dec 2023 07:07:25 +0000

Introduction

In the last article we made our datasets and learners more generic. Before writing this article, I realized we don't have any indication of the loss of our model during training, or the validation losses after training for the validation set. Let's add that capability in this article.

Saving weights is easy

So I want to save the weights of the model, and I'd like to do it often so that we can kill the program and pick up from where we left off fairly easily.

So I think I've decided to make an option on our Learner builder to save every batch. And while we're at it, so I don't have to update this code over multiple passes, we also need a way to check the validation loss of our learner. And it would be nice to log out the training loss each epoch, and even more, we can change the start epoch and load the old models so we can start from the middle of a training session. So let's just add all of this stuff! I'll comment about all the new stuff.

    pub struct Builder<'a, T, const N: usize, Category> {
        pub(super) device: AutoDevice,
        // Renamed this
        pub(super) train_dataset: Option<&'a DirectoryImageDataset<'a, N, Category>>,
        // Added validation set
        pub(super) valid_dataset: Option<&'a DirectoryImageDataset<'a, N, Category>>,
        pub(super) model: Option<Resnet34Model<N, f32>>,
        // Save each block
        pub(super) save_each_block: bool,
        // When should we start?
        pub(super) start_epoch: usize,
        pub(super) _phantom: PhantomData<T>,
    }

Then we can add an impl block to Builder that is generic over type T.

impl<'a, const N: usize, Category: IntoOneHot<N>, T> Builder<'a, T, N, Category> {
    // Should we save each block?
    pub fn save_each_block(mut self) -> Self {
        self.save_each_block = true;
        self
    }

    // Start epoch to builder
    pub fn start_epoch(mut self, start_epoch: usize) -> Self {
        self.start_epoch = start_epoch;
        self
    }

    // Add the validation dataset to the builder
    pub fn with_valid_dataset(
        mut self,
        valid_dataset: &'a DirectoryImageDataset<'a, N, Category>,
    ) -> Self {
        self.valid_dataset = Some(valid_dataset);
        self
    }
}

And now we can change the build() method to return the new structure of the VisualLearner.

    impl<'a, const N: usize, Category: IntoOneHot<N>> Builder<'a, Ready, N, Category> {
        pub fn build(self) -> VisualLearner<'a, N, Category> {
            let model = self.model.unwrap();
            VisualLearner::new(
                self.device,
                self.train_dataset.unwrap(),
                self.valid_dataset,
                self.model.unwrap(),
                self.save_each_block,
                self.start_epoch,
            )
        }
    }

And quickly modify the private constructor for VisualLearner

 pub struct VisualLearner<'a, const N: usize, Category> {
    device: AutoDevice,
    train_dataset: &'a DirectoryImageDataset<'a, N, Category>,
    // New, make this optional, so we don't always need to specify it.
    valid_dataset: Option<&'a DirectoryImageDataset<'a, N, Category>>,
    model: Resnet34Model<N, f32>,
    optimizer: Adam<Resnet34Built<N, f32>, f32, AutoDevice>,
    save_each_block: bool,
    start_epoch: usize,
}

const BATCH_SIZE: usize = 16;

impl<'a, const N: usize, Category: IntoOneHot<N>> VisualLearner<'a, N, Category> {
    // [snip builder]
   fn new(
        device: AutoDevice,
        train_dataset: &'a DirectoryImageDataset<'a, N, Category>,
        valid_dataset: Option<&'a DirectoryImageDataset<'a, N, Category>>,
        model: Resnet34Model<N, f32>,
        save_each_block: bool,
        start_epoch: usize,
    ) -> Self {
        let adam = Adam::new(&model.model, AdamConfig::default());
        Self {
            device,
            train_dataset,
            valid_dataset,
            model,
            optimizer: adam,
            save_each_block,
            start_epoch,
        }
    }
}

Now we just need to update the learn() method.

pub fn train(&mut self, epochs: usize) -> Result<(), Error> {
    let mut rng = rand::thread_rng();
    let mut grads = self.model.model.alloc_grads();
    let mut total_epoch_loss = 0.0;
    let mut num_batches = 0;
    let start = Instant::now();
    for epoch in self.start_epoch..self.start_epoch + epochs {
        log::info!("Epoch {}", epoch);
        for (image, is_cat) in self
            .train_dataset
            .shuffled(&mut rng)
            .map(Result::unwrap)
            .map(|(image, is_cat)| (image, is_cat.into_one_hot(&self.device)))
            .batch_exact(Const::<BATCH_SIZE>)
            .collate()
            .stack()
            .progress()
        {
            let logits = self.model.model.forward_mut(image.traced(grads));
            let loss = cross_entropy_with_logits_loss(logits, is_cat);
            total_epoch_loss += loss.array();
            num_batches += 1;

            grads = loss.backward();
            self.optimizer.update(&mut self.model.model, &grads)?;
            self.model.model.zero_grads(&mut grads);

            // Save the model after each block
            if self.save_each_block {
                self.save(format!("model-epoch-{}.safetensors", epoch))?;
            }
        }
        let dur = start.elapsed();

        // Log out the stats for the epoch
        log::info!(
            "Epoch {epoch} in {dur:?} ({:.3} batches/s): avg sample loss: {:.5}",
            num_batches as f32 / dur.as_secs_f32(),
            BATCH_SIZE as f32 * total_epoch_loss / num_batches as f32,
        );

        // If we have a valid set, go ahead and log out the validation loss.
        // Though we can remove this if it takes too long.
        if self.valid_dataset.is_some() {
            let valid_loss = self.valid_loss()?;
           log::info!("Valid loss: {:.5}", valid_loss);
        }
    }
    Ok(())
}

As you can see, saving the model each batch is a single call to the save() method. Lets define that and valid_loss() now.

impl<'a, const N: usize, Category: IntoOneHot<N>> VisualLearner<'a, N, Category> {
    pub fn valid_loss(&mut self) -> Result<f32, Error> {
        let mut total_epoch_loss = 0.0;
        let mut num_batches = 0;
        log::info!("Calculating validation loss");
        for (img, is_cat) in self
            .valid_dataset
            // If this is called without setting valid_dataset, it's an error
            .ok_or(Error::NoValidationDataset)?
            // NOT `shuffled()`, this just needs to iterate through once
            .iter()
            .map(Result::unwrap)
            .map(|(image, is_cat)| (image, is_cat.into_one_hot(&self.device)))
            .batch_exact(BATCH_SIZE)
            .collate()
            .stack()
            .progress()
        {
            let logits = self.model.model.forward(img);
            let loss = cross_entropy_with_logits_loss(logits, is_cat);
            total_epoch_loss += loss.array();
            num_batches += 1;
        }
        Ok(BATCH_SIZE as f32 * total_epoch_loss / num_batches as f32)
    }

    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), Error> {
        Ok(self.model.model.save_safetensors(path)?)
    }
}

Saving is easy, just a single call to save_safetensors(). The validation loss is a little more complicated. It's very similar to the training loop, but instead of calling shuffled(), we just call iter() because it doesn't matter if it is out of order, we just need to get through the validation set, and add the loss together.

The final step is to actually use this new stuff in the main.rs file. I'm going to make our chapter1 executable a little more complicated by adding some command line options. This is going to be pretty easy with the clap crate, and the derive feature flag.

#[derive(Debug, Parser)]
#[command(author = "Favil Orbedios")]
struct Args {
    /// The seed to create the [AutoDevice] with, default 0
    #[clap(long, short = 's', default_value = "0")]
    seed: u64,

    /// If set, load the model from this file
    #[clap(long = "model", short = 'm')]
    model_file: Option<PathBuf>,

    /// The epoch to start training at, default 0
    #[clap(long = "epoch", short = 'e', default_value = "0")]
    start_epoch: usize,

    /// The number of epochs to train for, default 3
    #[clap(long, short = 'n', default_value = "3")]
    epochs: usize,
}

This structure will be parsed from the command line arguments, so we can run it like

cargo run --release -- --seed 42 --epoch 10 --epochs 5 --model "model-epoch-9.safetensors"

And that will load the model from the model-epoch-9.safetensors file, start at epoch 10, train for 5 epochs, and use the seed 42 to initialize the AutoDevice. How do we use it now. I'll just paste the entire main() function we have so far.

fn main() -> Result<()> {
    env_logger::Builder::new()
        .filter_level(log::LevelFilter::Info)
        .init();
    color_eyre::install()?;

    // NEW: Parse the arguments from the command line
    let args = Args::parse();

    let path: PathBuf = untar_images(DatasetUrl::Pets)
        .context("downloading Pets")?
        .join("images");
    log::info!("Images are in: {}", path.display());

    // NEW: Set the seed directly, instead of using 0
    let dev = AutoDevice::seed_from_u64(args.seed);

    // Silly thing about the Pets dataset, all the cats have a capital first letter in their
    // filename, all the dogs are lowercase only
    let is_cat = |path: &Path| {
        path.file_name()
            .and_then(|n| n.to_str())
            .and_then(|n| n.chars().next().map(|c| c.is_uppercase()))
            .unwrap_or(false)
    };

    let dataset_loader = DirectoryImageDataLoader::builder(path, dev.clone())
        .with_label_fn(&is_cat)
        .with_splitter(RatioSplitter::with_seed_validation(0, 0.1))
        .build()?;
    let dataset = dataset_loader.training();
    log::info!("Found {} files", dataset.files().len());

    log::info!("Building the ResNet-34 model");
    let mut model = Resnet34Model::<2, f32>::build(dev.clone());
    log::info!("Done building model");

    // NEW: If we set the model file, let's load it
    if let Some(model_file) = args.model_file {
        log::info!("Loading old model");
        model.load_model(model_file)?;
        log::info!("Done loading old model");
    }

    let mut learner = VisualLearner::builder(dev.clone())
        .save_each_block()
        // NEW: starting epoch
        .start_epoch(args.start_epoch)
        .with_valid_dataset(dataset_loader.validation())
        .with_train_dataset(dataset)
        .with_model(model)
        .build();

    let valid_loss = learner.valid_loss()?;
    log::info!("Valid loss: {:.5}", valid_loss);

    log::info!("Training");
    // NEW: train for the specified number of epochs
    learner.train(args.epochs)?;
    log::info!("Done training");

    learner.save("model.safetensors")?;

    Ok(())
}

Run times

On my CPU bound laptop, this takes about 5 hours to run in release mode, with default arguments, because my GPU isn't supported yet. If we loaded it up on a Cuda supported computer, it would probably run a lot faster. I may spin up an instance on EC2 and see how the run times compare.

If anyone runs this code on a supported computer with the cuda feature flag, let me know in the comments how much faster it runs.

Conclusion

As always, I've uploaded the code to Github. You can check it out from the chapter-2-2 tag.

Tune in next time to learn about visualizing the images on the command line! I'm very excited to see how we do.

fast.ai Book in Rust - Chapter 2 - Part 1

Favil Orbedios — Wed, 29 Nov 2023 22:37:31 +0000

Introduction

In Part 6, we finished Chapter 1. We didn't get to do everything that we'd wanted to because my beefy desktop was put out of commission. That's still the case. Luckily, in this chapter, we don't need to focus too much on the training aspect.

This chapter focuses on defining the DataLoader classes and a Bing Image Search downloader that is provided with the fastai library. We're not going to implement a Bing downloader. That is too much work for something that could be a crate on its own. Please feel free to write such a crate, though, the world could use one.

One thing I'm noticing in the chapter however, that I clearly forgot about in the last chapter, is splitting the dataset into different sets, the training set, the validation set, and finally that test set. They are all important, and they all have separate uses. So I think now might be a good time to go through our code and refactor it to split out the different sets.

I will also take this opportunity to clean things up and make the code more generic.

Code Clean Up

The first step is moving tardyai/src/datasets.rs to a new module. Since this is specifically a category based module that is categorizing the input some how, I created the category module, and moved the file to tardyai/src/category/datasets.rs.

The next module I created was for the labels. In tardyai/src/category/encoders.rs, because it is for encoding a category into a one-hot tensor like we did for is_cat

pub type LabelFn<Category> = dyn Fn(&Path) -> Category;

This is the same function we were storing in our DirectoryImageDataset, I just changed it to use a generic return type instead of hard-coding it to bool. Though technically this isn't as generic as it could be, since we are assuming Paths. But for the foreseeable future, this is going to be coming directly from the hard drive. I don't have any plans to fetch inputs from a database or directly from the internet yet.

Now I want to be able to define a trait that will convert the Category into a one-hot encoded tensor. And since we are currently using bool, we can define the trait for bool while we're at it.

pub trait IntoOneHot<const N: usize>: Default {
    fn into_one_hot(&self, dev: &AutoDevice) -> Tensor<Rank1<N>, f32, AutoDevice>;
}

impl IntoOneHot<2> for bool {
    fn into_one_hot(&self, dev: &AutoDevice) -> Tensor<Rank1<2>, f32, AutoDevice> {
        let mut t = dev.zeros::<(Const<2>,)>();
        t[[0]] = !*self as usize as f32;
        t[[1]] = *self as usize as f32;
        t
    }
}

Now we need to create the idea of a DataLoader. The data loader in the fastai is able to split the input data into the different training, validation and testing data sets. Let's define a DirectoryImageDataLoader

// We need to thread N and Category through our types
pub struct DirectoryImageDataLoader<'fun, const N: usize, Category> {
    training: DirectoryImageDataset<'fun, N, Category>,
    validation: DirectoryImageDataset<'fun, N, Category>,
    test: DirectoryImageDataset<'fun, N, Category>,
}

impl<'fun, const N: usize, Category: IntoOneHot<N>> 
    DirectoryImageDataLoader<'fun, N, Category> 
{
    // Builder pattern again, its just so good
    pub fn builder(
        parent: impl AsRef<Path>,
        dev: AutoDevice,
    ) -> image_data_loader::Builder<'fun, N, Category> {
        image_data_loader::Builder::new(parent.as_ref().to_owned(), dev)
    }

    // Accessor methods
    pub fn training(&self) -> &DirectoryImageDataset<'fun, N, Category> {
        &self.training
    }

    pub fn validation(&self) -> &DirectoryImageDataset<'fun, N, Category> {
        &self.validation
    }

    pub fn test(&self) -> &DirectoryImageDataset<'fun, N, Category> {
        &self.test
    }
}

You should be able to see how I'm using a builder there, let's look at how that was made.

pub mod image_data_loader {
    // [snip use statements] 

    pub struct Builder<'fun, const N: usize, Category> {
        parent: PathBuf,
        dev: AutoDevice,
        splitter: Option<Box<dyn Splitter<PathBuf>>>,
        label_fn: Option<&'fun LabelFn<Category>>,
    }

    impl<'fun, const N: usize, Category: IntoOneHot<N>> Builder<'fun, N, Category> {
        pub fn new(parent: PathBuf, dev: AutoDevice) -> Self {
            Self {
                parent,
                dev,
                splitter: None,
                label_fn: None,
            }
        }

        // I'll talk about Splitter soon, I promise
        pub fn with_splitter(mut self, splitter: impl Splitter<PathBuf> + 'static) -> Self {
            self.splitter = Some(Box::new(splitter));
            self
        }

        pub fn with_label_fn(mut self, label_fn: &'fun LabelFn<Category>) -> Self {
            self.label_fn = Some(label_fn);
            self
        }

        pub fn build(self) -> Result<DirectoryImageDataLoader<'fun, N, Category>, Error> {
            let exts = image_extensions();

            let mut splitter = self
                .splitter
                .unwrap_or_else(|| Box::new(RatioSplitter::default()));
            // By default the label will return the default value regardless of input
            // This isn't useful, but I didn't want to make label_fn a required argument
            let label_fn = self.label_fn.unwrap_or(&|_| Default::default());

            // Walk the directory here, instead of the Dataset constructor
            // so we can split the datasets out
            let walker = WalkDir::new(self.parent).follow_links(true).into_iter();
            let files: Vec<_> = walker
                .filter_map(|entry| {
                    let entry = entry.ok()?;
                    entry
                        .path()
                        .extension()
                        .and_then(|ext| Some(exts.contains(ext.to_str()?)))?
                        .then_some(entry)
                })
                .map(|entry| entry.path().to_owned())
                .collect();
            let (training, validation, test) = splitter.split(files);
            let training = DirectoryImageDataset::new(
                &training, self.dev.clone(), label_fn
            )?;
            let validation = DirectoryImageDataset::new(
                &validation, self.dev.clone(), label_fn
            );
            let test = DirectoryImageDataset::new(&test, self.dev, label_fn)?;

            Ok(DirectoryImageDataLoader {
                training,
                validation,
                test,
            })
        }
    }
}

This walks the directory, and runs the Splitter on the files. Then constructs the datasets for use by the user.

Now let's talk about that Splitter, like I promised.

pub trait Splitter<T> {
    fn split(&mut self, files: Vec<T>) -> (Vec<T>, Vec<T>, Vec<T>);
}

I've gone ahead and allowed it to be generic over the type of input, since this is super straightforward to implement. The first implementation of a Splitter we'll create is the RatioSplitter. This will take in default ratios for the sizes of the validation and test datasets, and randomly (but deterministically) divide the input list into the various datasets.

pub struct RatioSplitter {
    rng: rand::rngs::StdRng,
    validation: f32,
    test: f32,
}

impl RatioSplitter {
    pub fn with_seed_validation_test(seed: u64, validation: f32, test: f32) -> Self {
        assert!(validation + test < 1.0);
        assert!(validation >= 0.0);
        assert!(test >= 0.0);
        let rng = rand::SeedableRng::seed_from_u64(seed);
        Self {
            rng,
            validation,
            test,
        }
    }

    pub fn with_seed_validation(seed: u64, validation: f32) -> Self {
        Self::with_seed_validation_test(seed, validation, 0.0)
    }

    // By default I'm going to assume your test data is somewhere else.
    pub fn with_seed(seed: u64) -> Self {
        Self::with_seed_validation_test(seed, 0.2, 0.0)
    }
}

// Very secure default seed, but it's good for reproducibility, and it's what dfdx uses
impl Default for RatioSplitter {
    fn default() -> Self {
        Self::with_seed(0)
    }
}

//       v--- Ensure that our types are sortable, for deterministic ordering
impl<T: Ord> Splitter<T> for RatioSplitter {
    fn split(&mut self, mut files: Vec<T>) -> (Vec<T>, Vec<T>, Vec<T>) {
        // Sort the files for deterministic ordering
        files.sort();
        // Shuffle with our deterministic Rng
        files.shuffle(&mut self.rng);

        let validation = (files.len() as f32 * self.validation) as usize;
        let test = (files.len() as f32 * self.test) as usize;

        let validation: Vec<T> = files.drain(..validation).collect();
        let test: Vec<T> = files.drain(..test).collect();
        let training: Vec<T> = files;

        (training, validation, test)
    }
}

Now we just need to change the DirectoryDataset, which I've also renamed to DirectoryImageDataset, since this only works for images.

pub struct DirectoryImageDataset<'fun, const N: usize, Category> {
    files: Vec<PathBuf>,
    dev: AutoDevice,
    label_fn: &'fun LabelFn<Category>,
    tensors: DashMap<PathBuf, Tensor<Rank3<3, 224, 224>, f32, AutoDevice>>,
}

impl<'fun, const N: usize, Category> DirectoryImageDataset<'fun, N, Category> {
    fn new(
        files: &[PathBuf],
        dev: AutoDevice,
        label_fn: &'fun LabelFn<Category>,
    ) -> Result<Self, Error> {
        Ok(Self {
            files: files.to_owned(),
            dev,
            label_fn,
            tensors: Default::default(),
        })
    }

    pub fn files(&self) -> &[PathBuf] {
        &self.files
    }
}

And finally we just need to update the call site in chapter1/src/main.rs.

    let is_cat = |path: &Path| {
        path.file_name()
            .and_then(|n| n.to_str())
            .and_then(|n| n.chars().next().map(|c| c.is_uppercase()))
            .unwrap_or(false)
    };

    let dataset_loader = DirectoryImageDataLoader::builder(path, dev.clone())
        .with_label_fn(&is_cat)
        .build()?;
    let dataset = dataset_loader.training();
    log::info!("Found {} files", dataset.files().len());

This code is just going to use the default Splitter, so we will see 20% of the files go to the validation set, and the final 80% go to the training set.

Conclusion

We've refactored out a DataLoader and cleaned up the interface to support labels other than bool. So that's cool.

As per the usual, the code is on Github, or you can checkout the tag chapter-2-1

Tune in next time and we are going to go over visualizing the datasets like fastai can do in jupyter. I've found a neat way to show images directly in the terminal.

Working through the fast.ai book in Rust - Part 6

Favil Orbedios — Tue, 28 Nov 2023 17:47:25 +0000

Introduction

In Part 5, we failed to correctly load weights into our model.

In this part, I will be talking about what we need to write to train our model, as if we were doing it from scratch. Eventually we will get the weights loaded, and this will be less time consuming, but for now, we need a training loop.

The Training Loop

Here is the code I want to be able to write to actually train our model.

    let model = Resnet34Model::<2, f32>::build(dev.clone());
    // model.download_model();

    let mut learner = VisualLearner::builder(dev.clone())
        .dataset(dataset)
        .model(model)
        .build();
    learner.train(10)?;

This looks pretty close to the python code that we're probably pretty familiar with at this point.

dls = ImageDataLoaders.from_name_func(
    path, get_image_files(path), valid_pct=0.2, seed=42,
    label_func=is_cat, item_tfms=Resize(224))

learn = vision_learner(dls, resnet34, metrics=error_rate)
learn.fine_tune(1)

But since this is Rust, I decided to use the Builder pattern to construct the learner. This allows us to add required arguments, and optional arguments in a type safe and easy to read way.

So how do we go about creating the learner? Let's dig into the code.

I first created a new module in tardyai/src/learners/visual.rs.

pub struct VisualLearner<'a> {
    device: AutoDevice,
    dataset: DirectoryImageDataset<'a>,
    model: Resnet34Model<2, f32>,
    optimizer: Adam<Resnet34Built<2, f32>, f32, AutoDevice>,
}

I'm hard-coding all the model specifics here for now. We can generalize this later. But for now, we have the device, the dataset, the model, and something called an optimizer. The optimizer is what takes the model, and calculates the changes to the weights that need to happen in order to get closer to the desired output.

In this code we're just going to use the Adam optimizer, because that is what the fast.ai book does.

So let's look at how we construct one of these. We can't just build it up from scratch in our main function, that wouldn't look good, and we don't want to necessarily construct all the options if we're going to just be using the default. In this particular model, I'm forcing the optimizer to be a default value, but we will be changing that in the future.

impl<'a> VisualLearner<'a> {
    pub fn builder(device: AutoDevice) -> builder::Builder<'a, builder::WithoutDataset> {
        builder::Builder {
            device,
            dataset: None,
            model: None,
            _phantom: PhantomData,
        }
    }
}

that creates a Builder with some None values and some PhantomData type. Let's look at how Builder is defined now to see what's going on. I always like to put my builders inside their own internal module, in order to separate out the namespace, and keep the implementation cleaner.

pub mod builder {
    use super::*;
    use crate::{datasets::DirectoryImageDataset, models::resnet::Resnet34Model};

    // 2 ZSTs
    pub struct WithoutDataset;
    pub struct WithoutModel;
    pub struct Ready;

    pub struct Builder<'a, T> {
        pub(super) device: AutoDevice,
        pub(super) dataset: Option<DirectoryImageDataset<'a>>,
        pub(super) model: Option<Resnet34Model<2, f32>>,
        // 1 PhantomData for typesafe builder pattern
        pub(super) _phantom: PhantomData<T>,
    }

    // 3 The first type of builder we get, can only add dataset
    impl<'a> Builder<'a, WithoutDataset> {
        pub fn dataset(
            self, 
            dataset: DirectoryImageDataset<'a>
        ) -> Builder<'a, WithoutModel> {
            Builder {
                device: self.device,
                dataset: Some(dataset),
                model: None,
                _phantom: PhantomData,
            }
        }
    }

    // 4 The second builder we get, still needs a model
    impl<'a> Builder<'a, WithoutModel> {
        pub fn model(self, model: Resnet34Model<2, f32>) -> Builder<'a, Ready> {
            Builder {
                device: self.device,
                dataset: self.dataset,
                model: Some(model),
                _phantom: PhantomData,
            }
        }
    }

    // 5 Finally construct the `VisualLearner`
    impl<'a> Builder<'a, Ready> {
        pub fn build(self) -> VisualLearner<'a> {
            let model = self.model.unwrap();
            VisualLearner::new(self.device, self.dataset.unwrap(), model)
        }
    }
}

So there is a lot going on here. Why do we have this PhantomData(1) thing with some weird T generic type tied to it? Well, that is actually how we can create builders that can't call build() until they are completely specified. In this case, we want the dataset and model to be specified before we allow the user to actually build the VisualLearner.

That is where those zero sized types(2) come into play. They are used as marker types to keep track of what we're allowed to do with our builder. WithoutDataset implies that the builder is missing it's dataset. WithoutModel clearly means the same about the model. And we have Ready which typically would have all the methods to add optional arguments to the builder.

In (3) we are defining the first method for the empty builder. Allowing us to set the dataset. In (4) we do the same for the model, and finally in (5) we construct the VisualLearner itself with a private constructor that we need to write.

Here is that constructor.

impl<'a> VisualLearner<'a> {
    // [snip]

    fn new(
        device: AutoDevice,
        dataset: DirectoryImageDataset<'a>,
        model: Resnet34Model<2, f32>,
    ) -> Self {
        let adam = Adam::new(&model.model, AdamConfig::default());
        Self {
            device,
            dataset,
            model,
            optimizer: adam,
        }
    }
}

This code takes the arguments supplied by the builder, constructs the optimizer, and returns a fully specified VisualLearner.

Now we've got the learner, how do we train it?

impl<'a> VisualLearner<'a> {
    // [snip]

    pub fn train(&mut self, epochs: usize) -> Result<(), Error> {
        let mut rng = rand::thread_rng();
        // 1 Allocate space for gradients in the device
        let mut grads = self.model.model.alloc_grads();
        for epoch in 0..epochs {
            log::info!("Epoch {}", epoch);
            for (image, is_cat) in self
                .dataset
                .shuffled(&mut rng) // 2 Shuffle the dataset randomly each iteration
                .map(Result::unwrap)
                .map(|(image, is_cat)| { // 3 convert the label into the One-hotted
                                         // representation that these models return
                    let mut one_hotted = [0.0; 2]; 
                    one_hotted[is_cat as usize] = 1.0;
                    (image, self.device.tensor(one_hotted))
                })
                .batch_exact(Const::<16>) // 4 Split it into batches of 16 images each
                .collate()                // 5 Split the tuples into two lists
                .stack()                  // 6 Combine lists to single batch tensors 
            {
                // 8              |--- 8 apply the model to the image batch
                //                v                        v- 7 Trace the gradients
                let logits = self.model.model.forward_mut(image.traced(grads));
                // 9 Calculate the difference between the output and the label
                let loss = cross_entropy_with_logits_loss(logits, is_cat);
                // 10 Calculate the gradients
                grads = loss.backward();
                // 11 Apply the optimizer
                self.optimizer.update(&mut self.model.model, &grads)?;
                // 12 Zero the gradients for next loop
                self.model.model.zero_grads(&mut grads);
            }
        }
        Ok(())
    }
}

First (1) thing a trainer needs to do is allocate room for the gradients. The gradients are determined by the slope of the many dimensional graph that is constructed by our giant tensor equation, and how large the loss is. So dfdx is calculating the derivative of all of the operations we're performing on our tensors, and applying the updates backward through the model.

Next (2) we're shuffling the dataset so we aren't over fitting to the same order of models. Then (3) we need to convert the label from a boolean, into a one-hotted tensor where the index for true or false (1 or 0) are set to 1, and all the others are zero.

In (4) we split the input into batches of 16 images, to make the training go faster. Then (5) collate the tuples into two side by side lists. And finally (6) stack those lists into tensors that are now 4 dimensional.

Inside our loop we (7) trace our input with our gradients, which just lets dfdx know that it needs to keep track of the differences so it can calculate the derivative. (8) Apply the model to the batch of images.

The (9) loss function figures out how far off we are from the actual proper label. This is used to (10) calculate the derivatives. We are trying to minimize this as much as possible, and it is what drives our optimizer (11). The final step in all of this is to zero out the gradients so we aren't accumulating previous runs.

Whew, that was a lot, but now we are done, and our code works! Unfortunately for me, my desktop with Cuda support is out of commission while I wait for new parts to arrive, and the laptop I'm using to write this article doesn't have an NVidia card. So until dfdx can support WebGPU, I'm stuck running it with a CPU Device, which is remarkably slow. So I can't train this model, but I invite anyone reading this to try it out and let me know how it goes in the comments.

Progress bars and stuff

One thing our library has been missing that fastai provides by default is an indication of progress. Progress while downloading models, while training, through the epochs. Let's see about adding that to our library.

To do this, we are going to use the indicatif crate. It provides progress bars, spinners and the like, and easy ways to update them.

For adding a progress bar to our training it's just a two line change. Add the indicatif::ProgressIndicator trait.

    for (image, is_cat) in self
        .dataset
        .shuffled(&mut rng)
        .map(Result::unwrap)
        .map(|(image, is_cat)| {
            let mut one_hotted = [0.0; 2];
            one_hotted[is_cat as usize] = 1.0;
            (image, self.device.tensor(one_hotted))
        })
        .batch_exact(Const::<16>)
        .collate()
        .stack()
        .progress() // <--- This is new
    {
        // [snip]
    }

Adding progress bars to our download is going to be only slightly trickier. In the download_file() function, we just need this.

    // [snip]
    log::info!("Downloading {} to: {}", &url, downloaded_file.display());
    let mut dest = File::create(&downloaded_file)?;
    let pb = indicatif::ProgressBar::new(response.content_length().unwrap_or(0));
    let mut buf = [0; 262144]; // 256KiB buffer
    while response.read(&mut buf)? > 0 {
        dest.write_all(&buf)?;
        pb.inc(buf.len() as u64);
    }
    Ok(downloaded_file)

That creates a progress bar with the length of the response as its length, and updates the progress bar every 256 KiB.

Conclusion

And that's it for this article. I added training to our model, and some fancy progress bars. Tune in next time. I think we can call Chapter 1 done, there was only the one project, and it's time to move on to the next thing.

As always, the code is available on Github or you can fetch the article-6 tag from the repo directly.

git pull
git checkout article-6

Stay tuned, and have a great day!

Working through the fast.ai book in Rust - Part 5

Favil Orbedios — Thu, 23 Nov 2023 19:09:43 +0000

Introduction

In Part 4, we managed to construct the model, and download the weights from Hugging Face.

In this part, I think I've figured out an ugly way to load those weights into the model. Let's see if we can do some classification with these weights and see if it gives something sensible at all.

Some cleanup first

So I'm going to admit that without the changes that are coming to dfdx in the github main branch, this is going to be remarkably hacky.

But before we get to that, I did some cleanup. While reading the fast.ai book, I stumbled across some terms that we were using incorrectly in our model. So in the interest of being true to the book, I've renamed Head and Tail to Stem and Head. In particular the head of the model is actually the last bit that performs the classification tasks. I figured it might be confusing to people who are reading this with a more ML background.

type Stem = (
    Conv2D<3, 64, 7, 2, 3>,
    BatchNorm2D<64>,
    ReLU,
    MaxPool2D<3, 2, 1>,
);

pub type Head<const NUM_CLASSES: usize> = (AvgPoolGlobal, Linear<512, NUM_CLASSES>);

I've also split Resnet18 and Resnet34 into two parts. The body, and a Head, which we've previously discussed. This will hopefully make it easier to do transfer learning and change the shape of the head later down the road. I've only copied the Resnet34 version here, since it's what we're working on.

// Layer clusters are in groups of [3, 4, 6, 4]
pub type Resnet34Body = (
    Stem,
    (
        BasicBlock<64>,
        ReLU,
        BasicBlock<64>,
        ReLU,
        BasicBlock<64>,
        ReLU,
    ),
    (
        (Downsample<64, 128>, ReLU, BasicBlock<128>, ReLU),
        (BasicBlock<128>, ReLU, BasicBlock<128>, ReLU),
    ),
    (
        (
            Downsample<128, 256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
        ),
        (
            BasicBlock<256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
        ),
    ),
    (
        Downsample<256, 512>,
        ReLU,
        BasicBlock<512>,
        ReLU,
        BasicBlock<512>,
        ReLU,
    ),
);

pub type Resnet34<const NUM_CLASSES: usize> = (Resnet34Body, Head<NUM_CLASSES>);

There, that looks better.

What about those weights?

Now we can start getting to the ugly stuff. The interface for visiting each of the tensors in a model, and updating them doesn't allow for a tree-like structure. It just iterates through the tensors one at a time, and passes the object you tell it to to your handler.

So I came up with the naive idea to just create a giant array of all of the names in the weights file, and manually sort them into the order that will be iterated over from the model we created. It's really not pretty, but it "works". (The scare quotes are there because it almost certainly is wrong somehow, and we'll have to wait until we try to classify something to see if these are sane values.)

I'm not going to paste the whole array, because it's > 200 lines, and that feels like a giant waste of space. If you want to go look at it, feel free to do so here.

pub(crate) const RESNET34_LAYERS: [&'static str; 254] = [
        "resnet.embedder.embedder.convolution.weight",
        "resnet.embedder.embedder.normalization.weight",
        "resnet.embedder.embedder.normalization.bias",
        "resnet.embedder.embedder.normalization.running_mean",
        "resnet.embedder.embedder.normalization.running_var",
        "resnet.embedder.embedder.normalization.num_batches_tracked",
        "resnet.embedder.embedder.normalization.num_batches_tracked",
        // [SNIP]
];

This snippet includes a part that I wanted to point out. The dfdx BatchNorm2D type contains two scalar values, epsilon and momentum, but the safetensors file downloaded from Hugging Face does not. It does have one scalar per normalization layer called num_batches_tracked. So in the interest of expediency, I just duplicated this value. This is most certainly not the correct thing to do, but the values I saw during my debugging process were all zeros, so I'm hoping it doesn't matter in the long run.

Now let's see about the class that takes that and applies the tensors to the model.

pub(crate) struct NamedTensorVisitor<'a> {
    names: &'a '[&'static str],
    idx: usize,
    tensors: &'a SafeTensors<'a>,
}

impl<'a> NamedTensorVisitor<'a> {
    pub(crate) fn new(names: &'a [&'static str], tensors: &'a SafeTensors<'a>) -> Self {
        Self {
            names,
            idx: 0,
            tensors,
        }
    }
}

So this class is dead simple, just increase the index every time you visit a tensor or scalar, and use the name that is at that index, fetching the weights from the SafeTensors field.

Implementing the TensorVisitor trait requires us to add the num-traits crate to our package, because we need to reference the NumCast trait in the definition.

impl<'a, E: Dtype + SafeDtype, D: Device<E>> TensorVisitor<E, D> for NamedTensorVisitor<'a> {
    type Viewer = ViewTensorMut;
    type Err = Error;
    type E2 = E;
    type D2 = D;

    fn visit<S: Shape>(
        &mut self,
        _opts: TensorOptions<S, E, D>,
        t: <Self::Viewer as TensorViewer>::View<'_, Tensor<S, E, D>>,
    ) -> Result<Option<Tensor<S, Self::E2, Self::D2>>, Self::Err> {
        log::debug!(
            "Loading tensor shape: {:?}, {:?}",
            t.shape(),
            &self.names.get(self.idx)
        );
        t.load_safetensor(
            &mut self.tensors,
            &self.names.get(self.idx).ok_or(Error::NotEnoughNames)?,
        )?;
        self.idx += 1;
        Ok(None)
    }

    fn visit_scalar<N: NumCast>(
        &mut self,
        _opts: ScalarOptions<N>,
        n: <Self::Viewer as TensorViewer>::View<'_, N>,
    ) -> Result<Option<N>, Self::Err> {
        log::debug!("Loading scalar: {:?}", &self.names.get(self.idx));
        let tensor = self
            .tensors
            .tensor(self.names.get(self.idx).ok_or(Error::NotEnoughNames)?)?;
        let data = tensor.data();
        let mut array = [0; 8];
        array.copy_from_slice(data);
        let val = f64::from_le_bytes(array);
        *n = N::from(val).ok_or(Error::NumberFormatException)?;

        self.idx += 1;
        Ok(None)
    }
}

This trait is a little complicated. We'll start by going over the associated types.

Viewer there is just assigned to ViewTensorMut which is a marker type that means we will receive a &mut Tensor<> in the visit method.

Err is just our crate's Error enumeration.

E2 and D2 are there in case you want to make changes to the tensors, either to the datatype or moving it to a different device. Since we aren't returning anything, this could be any number of things, but I'll just keep it set to the same as the inputs.

visit() is pretty straightforward, just load the tensor from the current name, and advance the index by one.

visit_scalar() was an interesting one, and I'm copying from the dfdx crate's implementation. Scalars in SafeTensors objects are stored as f64 in little-endian format, but the only way we can get to the data is by copying it into a [u8; 8], and converting that to an f64. Finally, we store it in the scalar by converting it to N(f32 in our case) from the f64, and advance the index by one.

Now we just have to actually use this struct to load up the tensors in the download_models method we created in the previous part.

    pub fn download_model(&mut self) -> Result<(), Error> {
        log::info!("Downloading model from {}", ModelUrl::Resnet34.url());
        let model_file = download_model(ModelUrl::Resnet34)?;

        // TODO: Somehow make something like this work
        // self.model.load_safetensors(&model_file)?;

        let file = File::open(model_file).unwrap();
        let buffer = unsafe { MmapOptions::new().map(&file).unwrap() };
        let tensors = SafeTensors::deserialize(&buffer).unwrap();

        let _ = <<Resnet34<N> as BuildOnDevice<AutoDevice, E>>::Built as TensorCollection<
            E,
            AutoDevice,
        >>::iter_tensors(&mut RecursiveWalker {
            m: &mut self.model,
            f: &mut NamedTensorVisitor::new(&RESNET34_LAYERS, &tensors),
        })?;

        Ok(())
    }

That complicated generic type after we load the tensors from the file really just boils down to interpretting the Built type, that we store as our model instance, as a TensorCollection<>, and calling the iter_tensors() associated function on that type.

iter_tensors() takes a ModuleVisitor. I'm using the included RecursiveWalker that just recurses through the tensors in the model, and calls visit() and visit_scalar() on them.

Actually, let's go ahead and clean that complicated generic type up a bit.

type Resnet34Built<const NUM_CLASSES: usize, E> =
    <Resnet34<NUM_CLASSES> as BuildOnDevice<AutoDevice, E>>::Built;

This will let us clean up our model type definition.

pub struct Resnet34Model<const NUM_CLASSES: usize, E>
where
    E: Dtype + SafeDtype,
    Resnet34<NUM_CLASSES>: BuildOnDevice<AutoDevice, E>,
    AutoDevice: Device<E>,
{
    pub model: Resnet34Built<NUM_CLASSES, E>,
}

And then actually clean up the ugly iter_tensors() call

    let _ = <Resnet34Built<N, E> as TensorCollection<E, AutoDevice>>::iter_tensors(
        &mut RecursiveWalker {
            m: &mut self.model,
            f: &mut NamedTensorVisitor::new(&RESNET34_LAYERS, &tensors),
        },
    )?;

That looks better. And we finally have the weights loaded into the model!

What can we do with it?

Let's try to use the model as it is to classify some of our Pets images, and see if the tensor it returns correlates to the cat or dog classes in the ImageNet dataset collection.

To apply the model to an image just means we call forward() on the tensor from the image. So let's do just that.

    model.download_model()?;

    let (image, is_cat) = dataset.get(1)?;
    let categories = model.model.forward(image);

    log::info!("Is Cat? {}", is_cat);
    log::trace!("Categories: {:#?}", categories.array());

    let max_category = categories
        .softmax()
        .array()
        .into_iter()
        .map(ordered_float::OrderedFloat)
        .enumerate()
        .max_by_key(|t| t.1)
        .unwrap();
    log::info!("(Category, Weight): {:?}", max_category);

This fetches the 2nd image in our dataset which, for me at least, is a cat. Then we apply the model to the tensor.

I've added a log statement for the categories tensor that is returned. This is a Rank1<1000> dimensional tensor that is supposed to have the largest value in the index that corresponds to the category of object in the input image.

So let's figure out which that is. I apply softmax() to the categories, which takes the ugly output of the original categories tensor, and massages the values so that they all add up to 1, while keeping the largest value the same. Unfortunately, f32 in rust doesn't implement the Ord trait, so we can't just fetch the maximum value from the tensor directly, we have to convert it to a type that does implement Ord. That is where the ordered-float crate comes into play. As long as we don't have any infinities or NaNs in our data, we can trust a total ordering.

And now if we run this we get (89, OrderedFloat(0.0099451095)). Hmm, if the weight of the category is supposed to correlate to how certain the model is, 0.9% doesn't inspire me with confidence. Let's check what that category is supposed to represent.

I found this json file that is supposed to be the categories that the original ResNet models were trained on. If we look in there for index 89, we get... sulphur-crested_cockatoo...

Well, that's a little frustrating, I don't think that is a cat. Something tells me that our model may not have been set up correctly. I'll attempt to get it working between articles. But this has taken a long time, and we may as well just train the model from scratch to see if we even have anything worth saving.

I'm going to plan on the next article talking about using the model in a learner with an optimizer like the penultimate line of the python code we're targetting.

learn = vision_learner(dls, resnet34, metrics=error_rate)

If I get this model working with transfer learning in the meantime, I'll switch it up and talk about how I managed it.

Anyway, until then, the code is available on github as always and if you're following along in your own editor, you can checkout the article-5 tag.

git pull origin
git checkout article-5

Stay tuned, and I hope to see you in the next article with a learner in tow.

Working through the fast.ai book in Rust - Part 4

Favil Orbedios — Wed, 22 Nov 2023 01:22:14 +0000

Introduction

In Part 3, we covered creating Tensors of our images, and loading them up into the device we're performing our matrix multiplication on.

In this part, I want to go over actually constructing the ResNet-34 model that the fast.ai book uses.

This is more difficult than it seems, because in addition to loading the weights from the internet, we also need to figure out how to cut off the last few layers, and add on some new layers that will give us our new categories.

The original model was trained on 1000 different categories, but in this first chapter, we're just determining if it is a cat or not. And the output Tensor is size 2.

So join me on this journey of discovery.

Let's build the model

First off, we need to enable support for convolutions. Convolutions take a group of pixels, in like a square around a central pixel, and convolve them into a single value. This is a fun relatively deep concept in math. I found this 3blue1brown video to be very approachable and entertaining.

Convolutions in dfdx are, unfortunately, only available on the nightly rust compiler. So first things first, let's enable the nightly channel.

[toolchain]
channel = "nightly"

Now I'm going to cheat a little. The lovely creators of dfdx left an example of the ResNet-18 structure in their repo. This isn't exactly what I need for ResNet-34, but it's pretty close, and it gives us a great jumping-off point.

So let's get started by stealing that code. I'm going to make a simple change and add a new type definition for the Tail of the structure. I plan to use this to change the shape of the final layers for re-training the model on different classes.

use dfdx::prelude::*;

type BasicBlock<const C: usize> = Residual<(
    Conv2D<C, C, 3, 1, 1>,
    BatchNorm2D<C>,
    ReLU,
    Conv2D<C, C, 3, 1, 1>,
    BatchNorm2D<C>,
)>;

type Downsample<const C: usize, const D: usize> = GeneralizedResidual<
    (
        Conv2D<C, D, 3, 2, 1>,
        BatchNorm2D<D>,
        ReLU,
        Conv2D<D, D, 3, 1, 1>,
        BatchNorm2D<D>,
    ),
    (Conv2D<C, D, 1, 2, 0>, BatchNorm2D<D>),
>;

type Head = (
    Conv2D<3, 64, 7, 2, 3>,
    BatchNorm2D<64>,
    ReLU,
    MaxPool2D<3, 2, 1>,
);

pub type Tail<const NUM_CLASSES: usize> = (AvgPoolGlobal, Linear<512, NUM_CLASSES>);

pub type Resnet18<const NUM_CLASSES: usize> = (
    Head,
    (BasicBlock<64>, ReLU, BasicBlock<64>, ReLU),
    (Downsample<64, 128>, ReLU, BasicBlock<128>, ReLU),
    (Downsample<128, 256>, ReLU, BasicBlock<256>, ReLU),
    (Downsample<256, 512>, ReLU, BasicBlock<512>, ReLU),
    Tail<NUM_CLASSES>,
);

Now, if you're anything like me, this is a bunch of incomprehensible nonsense. But luckily, we don't have to worry about why all these things exist, we just need to understand the basic structure.

For understanding what everything does, I found that Chapter 14 of the fast.ai book contains a wealth of information about the exact structure of what goes into a ResNet model. But that doesn't help us much right now, we're only on Chapter 1!

So in my search for how to understand the vague reasons for the large parts of this, and in particular, how to change this to ResNet-34, and even larger models; I went searching for the weights for the ResNet models, and I landed on the Hugging Face website. This is an awesome resource, and it definitely has the weights we need, but it also had a diagram of the structure of the layers of this model in particular. In the model card of the ResNet-34 model, I found this diagram.

Well, I was able to connect the dots, and I noticed the Residual words in the ResNet-18 types. The top diagram in that image is almost describing the BasicBlock and Downsample parts of this code. The 4 tuples between Head and Tail line up almost perfectly with those colorful blocks.

So BasicBlock corresponds to the first image. It's a pair of convolution layers, evidently separated by a ReLU layer, whatever that is. And Downsample must correspond to the second image with the dashed line. It looks like it takes the 64s to 128s, and that looks like what this line is doing.

    (Downsample<64, 128>, ReLU, BasicBlock<128>, ReLU),

Now, the difference seems to be just the number of BasicBlocks. ResNet-18 has 4 groups with 2 blocks in each group. But ResNet-34 seems to have 4 groups, with 3 in the first, then 4, then 6, then finally 3.

I still wasn't convinced, and I didn't trust my basic counting skills, so I went searching for how the fastai library defines the resnet34 model. That lead me to this code in the pytorch library. And sure enough, there are those numbers again!

def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                   **kwargs)

def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)

def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)

4 groups of 2, and [3, 4, 6, 3]. It even says BasicBlock! So it looks like the dfdx creators knew what they were doing. Who'd have thought?

So, it seems I just need to modify the structure of those middle tuples to have 3, 4, 6, and 3 layers. Let's see what we can do.

// Layer clusters are in groups of [3, 4, 6, 4]
pub type Resnet34<const NUM_CLASSES: usize> = (
    Head,
    (
        BasicBlock<64>,
        ReLU,
        BasicBlock<64>,
        ReLU,
        BasicBlock<64>,
        ReLU,
    ),
    (
        Downsample<64, 128>,
        ReLU,
        BasicBlock<128>,
        ReLU,
        BasicBlock<128>,
        ReLU,
        BasicBlock<128>,
        ReLU,
    ),
    (
        Downsample<128, 256>,
        ReLU,
        BasicBlock<256>,
        ReLU,
        BasicBlock<256>,
        ReLU,
        BasicBlock<256>,
        ReLU,
        BasicBlock<256>,
        ReLU,
        BasicBlock<256>,
        ReLU,
    ),
    (
        Downsample<256, 512>,
        ReLU,
        BasicBlock<512>,
        ReLU,
        BasicBlock<512>,
        ReLU,
    ),
    Tail<NUM_CLASSES>,
);

I can't tell, but this looks reasonable. Let's call it good for now.

You may have noticed that I included the ResNet-50 definition earlier. I wanted that there to point out that I have no idea what Bottleneck does, and at this point, I think it's fair to say it doesn't matter for Chapter 1, we only need the ResNet-34 model.

What can we do about the weights

Now Hugging Face has the weights we need. I've been reading up and it looks like the safetensors format is the fastest and safest weight format, and it just so happens to be supported by dfdx, we just need to enable the safetensors feature flag. Easy enough. Let's modify the top level Cargo.toml file.

dfdx = { version = "0.13", features = ["safetensors"] }

Now I've realized that I called my Url enum by a silly name. What kind of Url? Well let's change the name to DatasetUrl since it is used for downloading the dataset data. I just did a "Rename variable" operation in my IDE, and it took care of all the locations it was used for me.

So let's create a new enum in tardy/src/download.rs

const HF_BASE: &'static str = "https://huggingface.co/";

#[derive(Debug, Clone, Copy)]
pub(crate) enum ModelUrl {
    Resnet18,
    Resnet34,
}

impl ModelUrl {
    pub(crate) fn url(self) -> String {
        match self {
            ModelUrl::Resnet18 => {
                format!("{HF_BASE}microsoft/resnet-18/resolve/main/model.safetensors?download=true")
            }
            ModelUrl::Resnet34 => {
                format!("{HF_BASE}microsoft/resnet-34/resolve/main/model.safetensors?download=true")
            }
        }
    }
}

And now I want to create a new wrapper type that can hold onto a model, and provide methods for it, like download_model(). This will allow ease of use at the original call site. The format I'm aiming for is the following.

    let mut model = Resnet34Model::<1000, f32>::build(dev);
    model.download_model()?;

Now dfdx seems to have a very complicated type system, so this is going to look really awful. I'll go over the worst bits.

pub struct Resnet34Model<const NUM_CLASSES: usize, E>
where
    E: Dtype,
    Resnet34<NUM_CLASSES>: BuildOnDevice<AutoDevice, E>,
    AutoDevice: Device<E>,
{
    model: <Resnet34<NUM_CLASSES> as BuildOnDevice<AutoDevice, E>>::Built,
}

impl<E, const N: usize> Resnet34Model<N, E>
where
    E: Dtype,
    AutoDevice: Device<E>,
    Resnet34<N>: BuildOnDevice<AutoDevice, E>,
{
    pub fn build(dev: AutoDevice) -> Self {
        let model = dev.build_module::<Resnet34<N>, E>();
        Self { model }
    }
}

The most important part here is the line:

    model: <Resnet34<NUM_CLASSES> as BuildOnDevice<AutoDevice, E>>::Built,

This line is the difference between the type that we specified, and the concrete type that gets built by the AutoDevice::build_module() method. We can't just store a field with type Resnet34<NUM_CLASSES>, that isn't usable directly. In particular, it doesn't have any notion of the datatype that will be used for the model, whether it is f32 or f64. So we have to specify that Resnet34<NUM_CLASSES> implements the BuildOnDevice<> trait, and use the associated type, BuildOnDevice::Built.

We ensure that our type does implement BuildOnDevice with this line:

    Resnet34<NUM_CLASSES>: BuildOnDevice<AutoDevice, E>,

We next need to ensure that the device supports the datatype we are using with the line:

    AutoDevice: Device<E>,

The build() method now just takes in the Device we create in main.rs, and constructs the model with the AutoDevice::build_module method.

First source of frustration

Now, this looks like it should work. We've got a model that is relatively concise, and looks very similar to the ResNet-18 model.

So, why when I build this do I get the following horrendous error message?

error[E0277]: the trait bound `((dfdx::prelude::Conv2D<3, 64, 7, 2, 3>, 
dfdx::prelude::BatchNorm2D<64>, ReLU, MaxPool2D<3, 2, 1>), 
(dfdx::prelude::Residual<(dfdx::prelude::Conv2D<64, 64, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<64>, ReLU, dfdx::prelude::Conv2D<64, 64, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<64>)>, ReLU, dfdx::prelude::Residual<(dfdx::prelude::Conv2D<64, 
64, 3, 1, 1>, dfdx::prelude::BatchNorm2D<64>, ReLU, dfdx::prelude::Conv2D<64, 64, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<64>)>, ReLU, dfdx::prelude::Residual<(dfdx::prelude::Conv2D<64, 
64, 3, 1, 1>, dfdx::prelude::BatchNorm2D<64>, ReLU, dfdx::prelude::Conv2D<64, 64, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<64>)>, ReLU), (GeneralizedResidual<(dfdx::prelude::Conv2D<64, 
128, 3, 2, 1>, dfdx::prelude::BatchNorm2D<128>, ReLU, dfdx::prelude::Conv2D<128, 128, 3, 1, 
1>, dfdx::prelude::BatchNorm2D<128>), (dfdx::prelude::Conv2D<64, 128, 1, 2>, 
dfdx::prelude::BatchNorm2D<128>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<128, 128, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<128>, ReLU, dfdx::prelude::Conv2D<128, 128, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<128>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<128, 128, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<128>, ReLU, dfdx::prelude::Conv2D<128, 128, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<128>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<128, 128, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<128>, ReLU, dfdx::prelude::Conv2D<128, 128, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<128>)>, ReLU), (GeneralizedResidual<(dfdx::prelude::Conv2D<128, 
256, 3, 2, 1>, dfdx::prelude::BatchNorm2D<256>, ReLU, dfdx::prelude::Conv2D<256, 256, 3, 1, 
1>, dfdx::prelude::BatchNorm2D<256>), (dfdx::prelude::Conv2D<128, 256, 1, 2>, 
dfdx::prelude::BatchNorm2D<256>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>, ReLU, dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>, ReLU, dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>, ReLU, dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>, ReLU, dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>, ReLU, dfdx::prelude::Conv2D<256, 256, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<256>)>, ReLU), (GeneralizedResidual<(dfdx::prelude::Conv2D<256, 
512, 3, 2, 1>, dfdx::prelude::BatchNorm2D<512>, ReLU, dfdx::prelude::Conv2D<512, 512, 3, 1, 
1>, dfdx::prelude::BatchNorm2D<512>), (dfdx::prelude::Conv2D<256, 512, 1, 2>, 
dfdx::prelude::BatchNorm2D<512>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<512, 512, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<512>, ReLU, dfdx::prelude::Conv2D<512, 512, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<512>)>, ReLU, 
dfdx::prelude::Residual<(dfdx::prelude::Conv2D<512, 512, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<512>, ReLU, dfdx::prelude::Conv2D<512, 512, 3, 1, 1>, 
dfdx::prelude::BatchNorm2D<512>)>, ReLU), (AvgPoolGlobal, dfdx::prelude::Linear<512, _>)): 
BuildOnDevice<Cpu, _>` is not satisfied

Well our first clue to figure this monstrosity out comes from the little section at the end.

BuildOnDevice<Cpu, _>` is not satisfied

So, that means our giant model type, Resnet34 which just so happens to expand out to that awful tuple above, doesn't implement BuildOnDevice. Well, in the previous section we just stated that we needed it to do just that.

The next clue comes a few lines down from that

    = help: the following other types implement trait `BuildOnDevice<D, E>`:
              ()
              (M1,)
              (M1, M2)
              (M1, M2, M3)
              (M1, M2, M3, M4)
              (M1, M2, M3, M4, M5)
              (M1, M2, M3, M4, M5, M6)

Ahah! Evidently tuples of Models only implement BuildOnDevice for varieties up to 6-tuples. We have an 8-tuple, and a 12-tuple! So it looks like we just need to split our too large tuples down into smaller bite sized pieces.

Lets do just that. Here is the new model now.

pub type Resnet34<const NUM_CLASSES: usize> = (
    Head,
    (
        BasicBlock<64>,
        ReLU,
        BasicBlock<64>,
        ReLU,
        BasicBlock<64>,
        ReLU,
    ),
    (
        // tuples are only supported with up to 6 items in `dfdx`
        (Downsample<64, 128>, ReLU, BasicBlock<128>, ReLU),
        (BasicBlock<128>, ReLU, BasicBlock<128>, ReLU),
    ),
    (
        // tuples are only supported with up to 6 items in `dfdx`
        (
            Downsample<128, 256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
        ),
        (
            BasicBlock<256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
            BasicBlock<256>,
            ReLU,
        ),
    ),
    (
        Downsample<256, 512>,
        ReLU,
        BasicBlock<512>,
        ReLU,
        BasicBlock<512>,
        ReLU,
    ),
    Tail<NUM_CLASSES>,
);

And that builds! Excellent it was a simple fix.

Moving on to downloading models

Now that we have a model that can be concretely represented, and the code builds, we need to add some code to download the model files from Hugging Face.

I'm going to go ahead and refactor tardyai/src/download.rs while I'm here, so we can reuse our old download logic.

// v--- refactor out this logic, so it's shorter in the other functions.
fn get_home_dir() -> Result<PathBuf, Error> {
    let home = homedir::get_my_home()?
        .expect("home directory needs to exist")
        .join(".tardyai");
    Ok(home)
}

pub fn untar_images(url: DatasetUrl) -> Result<PathBuf, Error> {
    let home = get_home_dir()?;
    let dest_dir = home.join("archive");
    ensure_dir(&dest_dir)?;
    let archive_file = download_file(url.url(), &dest_dir, None)?;

    let dest_dir = home.join("data");
    let dir = extract_archive(&archive_file, &dest_dir)?;

    Ok(dir)
}

// v--- Add a crate public function that will download from a `ModelUrl`
pub(crate) fn download_model(url: ModelUrl) -> Result<PathBuf, Error> {
    let home = get_home_dir()?;
    let dest_dir = home.join("models");
    ensure_dir(&dest_dir)?;
    let model_file = download_file(url.url(), &dest_dir, Some(&format!("{url:?}.safetensors")))?;
    Ok(model_file)
}

// v--- Change the name to something more generic
fn download_file(
    url: String,
    dest_dir: &Path,
    // v--- This was needed because the filenames we download from Hugging Face 
    //      are pretty ugly looking strings of hex digits.
    default_name: Option<&str>,
) -> Result<PathBuf, Error> {
    let mut response = reqwest::blocking::get(&url)?;
    let file_name = default_name
        .or(response.url().path_segments().and_then(|s| s.last()))
        .and_then(|name| if name.is_empty() { None } else { Some(name) })
        //     v--- Add a new `Error` variant
        .ok_or(Error::DownloadNameNotSpecified(url.clone()))?;

    let downloaded_file = dest_dir.join(file_name);

    // TODO: check if the archive is valid and exists
    if downloaded_file.exists() {
        log::info!("File already exists: {}", downloaded_file.display());
        return Ok(downloaded_file);
    }

    log::info!("Downloading {} to: {}", &url, downloaded_file.display());
    let mut dest = File::create(&downloaded_file)?;
    response.copy_to(&mut dest)?;
    Ok(downloaded_file)
}

And that's done, so let's create the download_models method on our concrete model type. We can even make it call load_safetensors() while we're at it.

impl<E, const N: usize> Resnet34Model<N, E>
where
    E: Dtype + dfdx::tensor::safetensors::SafeDtype,
    AutoDevice: Device<E>,
    Resnet34<N>: BuildOnDevice<AutoDevice, E>,
{
    // ...

    pub fn download_model(&mut self) -> Result<(), Error> {
        log::info!("Downloading model from {}", ModelUrl::Resnet34.url());
        let model_file = download_model(ModelUrl::Resnet34)?;
        self.model.load_safetensors(&model_file)?;
        Ok(())
    }
}

So that's it, we're done!

Not so fast

Ah, it builds fine now, but I'm now getting this error when it runs.

➜   cargo run
   Compiling tardyai v0.1.0 (/home/klah/git/articles/fastai-rust/tardyai/tardyai)
   Compiling chapter1 v0.1.0 (/home/klah/git/articles/fastai-rust/tardyai/chapter1)
    Finished dev [unoptimized + debuginfo] target(s) in 4.37s
     Running `target/debug/chapter1`
[2023-11-22T00:01:42Z INFO  tardyai::download] File already exists: /home/klah/.tardyai/archive/oxford-iiit-pet.tgz
[2023-11-22T00:01:42Z INFO  tardyai::download] Extracting archive /home/klah/.tardyai/archive/oxford-iiit-pet.tgz to: /home/klah/.tardyai/data
[2023-11-22T00:01:42Z INFO  tardyai::download] Archive already extracted to: /home/klah/.tardyai/data/oxford-iiit-pet/
[2023-11-22T00:01:42Z INFO  chapter1] Images are in: /home/klah/.tardyai/data/oxford-iiit-pet/images
[2023-11-22T00:01:42Z INFO  chapter1] Found 7390 files
[2023-11-22T00:01:48Z INFO  tardyai::models::resnet] Downloading model from https://huggingface.co/microsoft/resnet-34/resolve/main/model.safetensors?download=true
[2023-11-22T00:01:49Z INFO  tardyai::download] Downloading https://huggingface.co/microsoft/resnet-34/resolve/main/model.safetensors?download=true to: /home/klah/.tardyai/models/Resnet34.safetensors
Error:
   0: Error with safetensors file: SafeTensorError(TensorNotFound("0.0.weight"))

Location:
   chapter1/src/main.rs:38

Backtrace omitted. Run with RUST_BACKTRACE=1 environment variable to display it.
Run with RUST_BACKTRACE=full to include source snippets.

Well, dang it. Evidently safetensors files have names for the layers that they are storing the weights for. I guess I'm going to have to figure out what this file actually contains, and load them individually into our model.

dfdx supports safetensors with the safetensors crate. So I'll add that dependency and lets get to debugging.

This page of the safetensors docs mentions using the memmap2 crate, so I'll go ahead and add that as well.

    pub fn download_model(&mut self) -> Result<(), Error> {
        log::info!("Downloading model from {}", ModelUrl::Resnet34.url());
        let model_file = download_model(ModelUrl::Resnet34)?;
        // self.model.load_safetensors(&model_file)?;

        let file = File::open(model_file).unwrap();
        let buffer = unsafe { MmapOptions::new().map(&file).unwrap() };
        let tensors = SafeTensors::deserialize(&buffer).unwrap();
        let mut names = tensors.tensors();
        names.sort_by_key(|t| t.0.clone());
        for (name, tensor) in names {
            log::info!("Name: {name}: {:?}", tensor.shape());
        }
        Ok(())
    }

Running this code gives us the following output.

[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: classifier.1.bias: [1000]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: classifier.1.weight: [1000, 512]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.embedder.embedder.convolution.weight: [64, 3, 7, 7]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.embedder.embedder.normalization.bias: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.embedder.embedder.normalization.num_batches_tracked: []
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.embedder.embedder.normalization.running_mean: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.embedder.embedder.normalization.running_var: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.embedder.embedder.normalization.weight: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.0.convolution.weight: [64, 64, 3, 3]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.0.normalization.bias: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.0.normalization.num_batches_tracked: []
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.0.normalization.running_mean: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.0.normalization.running_var: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.0.normalization.weight: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.1.convolution.weight: [64, 64, 3, 3]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.1.normalization.bias: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.1.normalization.num_batches_tracked: []
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.1.normalization.running_mean: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.1.normalization.running_var: [64]
[2023-11-22T00:25:47Z INFO  tardyai::models::resnet] Name: resnet.encoder.stages.0.layers.0.layer.1.normalization.weight: [64]
...

I've cut off most of the output, because it is an awful lot.

Well the shape of resnet.embedder.embedder.confolution.weight is very similar to the convolution in Head:

type Head = (
    Conv2D<3, 64, 7, 2, 3>,
    BatchNorm2D<64>,
    ReLU,
    MaxPool2D<3, 2, 1>,
);

So I bet they correlate, and this all looks fairly structured. stages runs from 0 to 3, so there are 4 stages, like we have 4 tuples of BasicBlocks.

Stage 2 has the first layer key running up to 5, so it has 6 layers, which is the number of BasicBlocks in the third tuple.

So I think we have the pattern pretty mapped out to what tensor goes where. Now we need to figure out how to actually load the tensors into the various weights, ideally without specifying each tuple entry manually.

This will probably be done with some form of TensorVisitor. But this article is getting pretty long, so let's save that for next time.

Conclusion

In Part 4, we were able to construct the model, and download the weights in the form of a safetensors file from Hugging Face. But we ran into an issue with actually loading the weights into the model, because they weren't named the same as what dfdx expects. Check out the code for this part at github. Or check out the article-4 tag.

git co article-4

Stay tuned for Part 5 where we figure out how to solve this conundrum.

Working through the fast.ai book in Rust - Part 3

Favil Orbedios — Tue, 21 Nov 2023 06:38:22 +0000

Introduction

In Part 2, we went over downloading and extracting our datasets from the internet.

Today in Part 3, we are going to go over reading that data and transforming it into a form that our models can actually understand.

Tensor objects are potentially multi-dimensional arrays that store numbers. And since files on a computer are all just numbers, we can utilize that to our advantage, and create something called a Rank 3 Tensor of our images. This is a 3 dimensional structure that represents a single image in it's entirety. The dimensions of the Tensor represent the height, width, and the color channels of the image.

Specifically for the ResNet model that chapter 1 wants us to work with these images need to be in the shape of 3x224x224. That corresponds to 3 color channels, and height and width both set to 224. Because of this size limitation, we are going to need to resize our images before we store them in our Tensors.

Refactoring

First things first, in the last Part of this series, we just wrote all of our code into a single module. That is going to get overwhelming quickly, so let's refactor our code into separate modules.

This is how our tardyai/src/lib.rs will look

pub mod download;
pub mod error;

pub use self::{
    download::{untar_images, Url},
    error::Error,
};

So in order to accomplish this, we can pull out the Error enum into tardyai/src/error.rs

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("reqwest error: {0}")]
    Reqwest(#[from] reqwest::Error),

    #[error("io error: {0}")]
    IO(#[from] std::io::Error),

    #[error("homedir error: {0}")]
    Home(#[from] homedir::GetHomeError),

    #[error("tar entry error: {0}")]
    TarEntry(&'static str),
}

And we can pull everything else out into tardyai/src/download.rs.

use std::{
    fs::File,
    io::{self, Seek},
    path::PathBuf,
};

use flate2::read::GzDecoder;
use tar::Archive;

use crate::error::Error;

const S3_BASE: &str = "https://s3.amazonaws.com/fast-ai-";
const S3_IMAGE: &str = "imageclas/";

#[derive(Debug, Clone, Copy)]
pub enum Url {
    Pets,
}

impl Url {
    pub fn url(self) -> String {
        match self {
            Self::Pets => {
                format!("{S3_BASE}{S3_IMAGE}oxford-iiit-pet.tgz")
            }
        }
    }
}

fn ensure_dir(path: &PathBuf) -> Result<(), Error> {
    if !path.exists() {
        std::fs::create_dir_all(path)?;
    }
    Ok(())
}

pub fn untar_images(url: Url) -> Result<PathBuf, Error> {
    let home = &homedir::get_my_home()?
        .expect("home directory needs to exist")
        .join(".tardyai");
    let dest_dir = home.join("archive");
    ensure_dir(&dest_dir)?;
    let archive_file = download_archive(url, &dest_dir)?;

    let dest_dir = home.join("data");
    let dir = extract_archive(&archive_file, &dest_dir)?;

    Ok(dir)
}

fn download_archive(url: Url, dest_dir: &PathBuf) -> Result<PathBuf, Error> {
    let mut response = reqwest::blocking::get(url.url())?;
    let archive_name = response
        .url()
        .path_segments()
        .and_then(|s| s.last())
        .and_then(|name| if name.is_empty() { None } else { Some(name) })
        .unwrap_or("tmp.tar.gz");

    let archive_file = dest_dir.join(archive_name);

    // TODO: check if the archive is valid and exists
    if archive_file.exists() {
        log::info!("Archive already exists: {}", archive_file.display());
        return Ok(archive_file);
    }

    log::info!(
        "Downloading {} to archive: {}",
        url.url(),
        archive_file.display()
    );
    let mut dest = File::create(&archive_file)?;
    response.copy_to(&mut dest)?;
    Ok(archive_file)
}

fn extract_archive(archive_file: &PathBuf, dest_dir: &PathBuf) -> Result<PathBuf, Error> {
    let tar_gz = File::open(archive_file)?;
    let tar = GzDecoder::new(tar_gz);
    let mut archive = Archive::new(tar);

    log::info!(
        "Extracting archive {} to: {}",
        archive_file.display(),
        dest_dir.display()
    );
    let dir = {
        let entry = &archive
            .entries()?
            .next()
            .ok_or(Error::TarEntry("No entries in archive"))??;
        entry.path()?.into_owned()
    };
    let archive_dir = dest_dir.join(dir);
    if archive_dir.exists() {
        log::info!("Archive already extracted to: {}", archive_dir.display());
        return Ok(archive_dir);
    }

    let tar = archive.into_inner();
    let mut tar_gz = tar.into_inner();
    tar_gz.seek(io::SeekFrom::Start(0))?;
    let tar = GzDecoder::new(tar_gz);
    let mut archive = Archive::new(tar);
    archive.unpack(dest_dir)?;

    Ok(archive_dir)
}

And that is done. Super simple, but much more organized.

On to finding images

Now that that's out of the way, we can get to the actual meat of this article.

As a good starting place let's create a new module for our datasets. And the associated tardyai/src/datasets.rs file.

pub mod datasets;

And now we need to create a struct that can keep track of all the image files we want to read.

use std::path::PathBuf;

pub struct DirectoryImageDataset {
    files: Vec<PathBuf>,
}

Let's also create a constructor that takes the parent directory, walks all the files, and collects all the images into the files field. This is going to use the walkdir crate for walking the parent directory and fetching the names of images that match extensions, and the image crate for listing the actual file extensions that are supported as well as eventually decoding the images and loading them into memory.

impl DirectoryImageDataset {
    pub fn new(parent: PathBuf) -> Result<Self, Error> {
        let exts = image_extensions();

        let walker = WalkDir::new(parent).follow_links(true).into_iter();
        let files = walker
            .filter_map(|entry| {
                let entry = entry.ok()?;
                entry
                    .path()
                    .extension()
                    .and_then(|ext| Some(exts.contains(ext.to_str()?)))?
                    .then_some(entry)
            })
            .map(|entry| entry.path().to_owned())
            .collect();
        Ok(Self { files })
    }
}

This uses a function called image_extensions that returns a HashSet of all the supported file extensions that I feel like including right now.

fn image_extensions() -> HashSet<&'static str> {
    let mut set = HashSet::default();
    set.extend(ImageFormat::Jpeg.extensions_str());
    set.extend(ImageFormat::Png.extensions_str());
    set.extend(ImageFormat::Gif.extensions_str());
    set.extend(ImageFormat::WebP.extensions_str());
    set.extend(ImageFormat::Tiff.extensions_str());
    set.extend(ImageFormat::Bmp.extensions_str());
    set.extend(ImageFormat::Qoi.extensions_str());
    set
}

I've also added a utility method for fetching the list of files that are represented

impl DirectoryImageDataset {
    // pub fn new(...)

    pub fn files(&self) -> &[PathBuf] {
        &self.files
    }
}

Let's get to the `Tensors` now

Now that we've discovered all the images, we just need to get them into memory. For now, I'm planning on doing this the naive way and just load the images each time they are requested. We can optimize it later by adding a cache

The dfdx crate has a trait called ExactSizeDataset that it can use to power the training epochs and calculate the validation loss at the end.

It's a pretty straightforward interface that has only two methods that are needed, a len(), and a get() method. len() is pretty self explanatory. get() takes an index and returns something you specify. Generally this can be both a Tensor, and the label that that Tensor represents.

impl ExactSizeDataset for DirectoryImageDataset<'_> {
    type Item<'a> = Result<(Tensor<Rank3<3, 224, 224>, f32, AutoDevice>, bool), Error>
    where
        Self: 'a;

    fn get(&self, index: usize) -> Self::Item<'_> {
        let image_file = &self.files[index];
        // Read the image and resize it to 224x224, and 3 channels
        let image = ImageReader::open(image_file)?
            .decode()?
            .resize_exact(224, 224, FilterType::Triangle)
            .into_rgb8();

        // Shrink the byte values to f32 between [0, 1]
        let bytes: Vec<f32> = image.as_bytes().iter().map(|&b| b as f32 / 255.0).collect();

        // Create the tensor and the label
        Ok((
            self.dev
                .tensor_from_vec(bytes, (Const::<3>, Const::<224>, Const::<224>)),
            (self.label_fn)(image_file),
        ))
    }

    fn len(&self) -> usize {
        self.files.len()
    }
}

I'm doing a lot here. Probably the most important is that I need to modify the struct to contain the device in order to actually construct the Tensor objects. This is because the Tensor might be stored on the GPU, and the device facilitates constructing them there.

I've also made the struct slightly generic, by storing a function that maps the path to the label, in this case the label is hard coded to a bool. We'll need to take care of that later, but I don't want to borrow tomorrow's problem just yet.

Here's the new structure for the struct, along with the associated constructor changes to support it. I unfortunately had to add a lifetime to handle the label function, but it's not too bad, we can get away with lifetime elision except when we are constructing it. I suppose I could have gone with 'static but that would have meant that I could only pass functions, not lambdas like I am hoping to eventually use. (So much for not borrowing tomorrow's problems, I guess)

pub struct DirectoryImageDataset<'fun> {
    files: Vec<PathBuf>,
    dev: AutoDevice,
    label_fn: Box<dyn Fn(&Path) -> bool + 'fun>,
}

impl<'fun> DirectoryImageDataset<'fun> {
    pub fn new(
        parent: PathBuf,
        dev: AutoDevice,
        label_fn: impl Fn(&Path) -> bool + 'fun,
    ) -> Result<Self, Error> {
        let exts = image_extensions();

        let walker = WalkDir::new(parent).follow_links(true).into_iter();
        let files = walker
            .filter_map(|entry| {
                let entry = entry.ok()?;
                entry
                    .path()
                    .extension()
                    .and_then(|ext| Some(exts.contains(ext.to_str()?)))?
                    .then_some(entry)
            })
            .map(|entry| entry.path().to_owned())
            .collect();
        Ok(Self {
            files,
            dev,
            label_fn: Box::new(label_fn),
        })
    }
}

How do we optimize it?

Right now, as we run through the epochs if we started training right now, we'd run into a problem with speed. We aren't caching the Tensors, so they are getting cleared and rebuilt each time through the training loop.

That will take much longer than if we just cached all the values the first time. So by spending a little memory (about 2 GiB from my back of the envelope math), we can improve our performance, significantly.

Now, the ExactSizeDataset doesn't allow us to get a mutable reference to self, so we can't just use a HashMap and be done, because all the points where we would be able to add tensors to the cache don't have a mutable reference.

So to help us fix that issue, and help us in the event that we are calling get() in parallel (say if we are trying to load all the images with rayon, for example), we can use the dashmap crate. DashMap is a concurrent hashmap, which means we can update the keys and values concurrently in multiple threads. This also means that we can insert values with only a shared reference.

pub struct DirectoryImageDataset<'fun> {
    files: Vec<PathBuf>,
    dev: AutoDevice,
    label_fn: Box<dyn Fn(&Path) -> bool + 'fun>,
    tensors: DashMap<PathBuf, Tensor<Rank3<3, 224, 224>, f32, AutoDevice>>,
}

We also need to set this in the constructor

// pub fn new(...) ... {
        // ...
        Ok(Self {
            files,
            dev,
            label_fn: Box::new(label_fn),
            tensors: Default::default(),
        })
}

And the final step is to pipe it through the get() method.

    fn get(&self, index: usize) -> Self::Item<'_> {
        let image_file = &self.files[index];

        // v---- New stuff here ---v
        let label = (self.label_fn)(image_file);
        if self.tensors.contains_key(image_file) {
            return Ok((self.tensors.get(image_file).unwrap().clone(), label));
        }
        // Read the image and resize it to 224x224, and 3 channels
        let image = ImageReader::open(image_file)?
            .decode()?
            .resize_exact(224, 224, FilterType::Triangle)
            .into_rgb8();

        // Shrink the byte values to f32 between [0, 1]
        let bytes: Vec<f32> = image.as_bytes().iter().map(|&b| b as f32 / 255.0).collect();

        // Create the tensor and the label
        let tensor = self
            .dev
            .tensor_from_vec(bytes, (Const::<3>, Const::<224>, Const::<224>));

        // v--- And here ---v
        self.tensors.insert(image_file.clone(), tensor.clone());
        Ok((tensor, label))
    }

Finally, I've updated the chapter1/src/main.rsfile to use the new struct.

    // ...
    let path: PathBuf = untar_images(Url::Pets)
        .context("downloading Pets")?
        .join("images");
    log::info!("Images are in: {}", path.display());

    // v--- New stuff ---v
    let dev = AutoDevice::default();

    // Silly thing about the Pets dataset, all the cats have a capital first letter in their
    // filename, all the dogs are lowercase only
    let is_cat = |path: &Path| {
        path.file_name()
            .and_then(|n| n.to_str())
            .and_then(|n| n.chars().next().map(|c| c.is_uppercase()))
            .unwrap_or(false)
    };

    let dataset = DirectoryImageDataset::new(path, dev.clone(), is_cat)?;
    log::info!("Found {} files", dataset.files().len());

    // ...

Conclusion

Well, that was more satisfying than the last two articles, we're finally actually able to store things on the GPU, or at least in a Tensor. To actually build this for GPU support, you need to build it with the cuda feature enabled.

Check out the code at github.

In Part 4, we're going to go through the process of building the ResNet-34 model in rust. Stay tuned! I'm so excited!

Working through the fast.ai book in Rust - Part 2

Favil Orbedios — Sun, 19 Nov 2023 21:07:08 +0000

Introduction

In Part 1 we introduced the dfdx crate. And we didn't get into any of the actually implementing any of the fast.ai book projects.

In Part 2 we are going to see how far we can get into chapter 1 of the book. Since this isn't python, and we don't have the fastai library, we are going to have to do everything ourselves.

If you want to follow along, and don't have a copy of the book, you can read it online for free here.

In particular this is what the book wants us to write:

from fastai.vision.all import *
path = untar_data(URLs.PETS)/'images'

def is_cat(x): return x[0].isupper()
dls = ImageDataLoaders.from_name_func(
    path, get_image_files(path), valid_pct=0.2, seed=42,
    label_func=is_cat, item_tfms=Resize(224))

learn = vision_learner(dls, resnet34, metrics=error_rate)
learn.fine_tune(1)

We can see from that, that it isn't much code. But the reason I don't like it, and the reason I'm writing this series is because it is just a bunch of magic. It gets you on your feet quickly, by hiding all the fun parts behind its façade.

In this sample we can see that it:

Automatically downloads, and extracts the images to an images folder.
Defines a label function
Automatically loads the images from the path with the ImageDataLoaders
Constructs a learner from a fully available resnet34 model with weights that are already downloaded.
And runs a learning algorithm on it for a single cycle.

Now this is too much to cover in a single article, so I'm going to focus on 1. for Part 2.

Creating a new Rust package

I realized while writing this article, that my structure for my code needs refinement. So I'm going to throw away the old code, and construct a repo with a number of crates in a Rust workspace.

If you want to follow along, I've created a git repo called tardyai, where I will be committing all my code to.

To fetch the specific tag from the repo use the following command:

git clone --branch START_HERE https://github.com/favilo/tardyai.git

That will download the repo and put you in the same starting point as me. Specifically, it contains a Rust workspace with two member crates: tardyai, and chapter1. Both of these are the default packages created by cargo new.

tardyai will be a small, incomplete port of the fastai library. It won't run any code itself, it just contains all the logic around downloading images, for now.

Let's add URLs

It would be very nice if we could take the same URLs that are in the python library and do the same thing in Rust.

I'm envisioning an interface similar the following. I'm adding this to our chapter1/src/main.rs file.

use std::path::PathBuf;

fn main() {
    let path: PathBuf = tardyai::untar_images(tardyai::Url::Pets)
        .join("images");
}

Now to make that a reality, lets edit tardyai/src/lib.rs

use std::path::PathBuf;

pub enum Url {
    Pets,
}

pub fn untar_images(url: Url) -> PathBuf {
    todo!()
}

This just panics, but at least everything compiles.

From here we need to convert that enum Url::Pets to an actual URL. For the fastai library this is https://s3.amazonaws.com/fast-ai-imageclas/oxford-iiit-pet.tgz. So lets add some methods to the Url type to get a URL.

const S3_BASE: &str = "https://s3.amazonaws.com/fast-ai-";
const S3_IMAGE: &str = "imageclas/";

// v-- I decided that we need to derive some sane traits by default.
#[derive(Debug, Clone, Copy)]
pub enum Url {
    Pets,
}

impl Url {
    pub fn url(self) -> String {
        match self {
            Self::Pets => {
                format!("{S3_BASE}{S3_IMAGE}oxford-iiit-pet.tgz")
            }
        }
    }
}

This defines the url() method, and I created a constant called S3_BASE in order to collect the common prefix. This will allow us to quickly add new paths, and their corresponding URLs.

Actually download something why don't you?

Now we need to actually connect to the internet and download our archive from S3. In order to do this I'm going to use the reqwest crate. This crate is the defacto crate for making HTTP requests. It offers both an async and a blocking API. We are going to be using the blocking API for now. (Maybe in a future article I'll convert everything over to async/await)

➜   cargo add reqwest -p tardyai -F blocking
    Updating crates.io index
      Adding reqwest v0.11.22 to dependencies.
             Features:
             + __tls
             + blocking
             + default-tls
             + hyper-tls
             + native-tls-crate
             + tokio-native-tls
             38 deactivated features

This adds the latest version of reqwest with the blocking feature turned on.

Then we edit tardyai/src/lib.rs

pub fn untar_images(url: Url) -> PathBuf {
    let response = reqwest::blocking::get(url.url()).expect("get failed");
    // ...
}

That .expect() looks pretty ugly. Let's clean that up with our own custom error type derived with the help of thiserror.

➜   cargo add -p tardyai thiserror
    Updating crates.io index
      Adding thiserror v1.0.50 to dependencies.

NOTE: I'm going to stop writing down the steps to add a crate. They are almost always the same. Instead I'll mention the crate and any features we need to add to get it to work for us.

thiserror will let us create an error type that is portable, and works with some nice error reporting crates that I'll talk about later.

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("reqwest error: {0}")]
    Reqwest(#[from] reqwest::Error),

    #[error("io error: {0}")]
    IO(#[from] std::io::Error),
}

Then we can change the signature of the untar_images.

pub fn untar_images(url: Url) -> Result<PathBuf, Error> {
    let response = reqwest::blocking::get(url.url())?;
    log::info!("response: {:?}", response);

    Ok(todo!())
}

So, that is us fetching the file from the URL. Of course this is useless to us as it stands, because we haven't saved it to the hard disk, but this will not use any bandwidth because we haven't fetched the body of the response.

Save it to the hard disk already

The fastai library fetches the archive files to ~/.fastai/archive/. I'm going to do the same thing, but in ~/.tardyai/archive/ instead.

So first we need to make sure that the directory exists. And we need to fetch the user's home in a cross platform manner. For that I'm using the homedir crate.

fn ensure_dir(path: &PathBuf) -> Result<(), Error> {
    if !path.exists() {
        std::fs::create_dir_all(path)?;
    }
    Ok(())
}

pub fn untar_images(url: Url) -> Result<PathBuf, Error> {
    let dest_dir = homedir::get_my_home()?
        .expect("home directory needs to exist")
        .join(".tardyai")
        .join("archive");
    ensure_dir(&dest_dir)?;
    // ...
}

This required creating a new variant for our Error enum. I called it Home.

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("reqwest error: {0}")]
    Reqwest(#[from] reqwest::Error),

    #[error("io error: {0}")]
    IO(#[from] std::io::Error),

    #[error("homedir error: {0}")]
    Home(#[from] homedir::GetHomeError),
}

And to save it to disk let's create a new function.

fn download_archive(url: Url, dest_dir: &PathBuf) -> Result<PathBuf, Error> {
    let mut response = reqwest::blocking::get(url.url())?;
    let archive_name = response
        .url()
        .path_segments()
        .and_then(|s| s.last())
        .and_then(|name| if name.is_empty() { None } else { Some(name) })
        .unwrap_or("tmp.tar.gz");

    let archive_file = dest_dir.join(archive_name);

    // TODO: check if the archive is valid and exists
    if archive_file.exists() {
        log::info!("Archive already exists: {}", archive_file.display());
        return Ok(archive_file);
    }

    log::info!(
        "Downloading {} to archive: {}",
        url.url(),
        archive_file.display()
    );
    let mut dest = File::create(&archive_file)?;
    response.copy_to(&mut dest)?;
    Ok(archive_file)
}

We have the archive, now what?

Well, let's decompress and extract it of course. For decompression I'm going to use the flate2 crate, with the rust_backend feature. And for extracting the resulting tar file, I'll use the tar crate.

fn extract_archive(archive_file: &PathBuf, dest_dir: &PathBuf) -> Result<(), Error> {
    let tar_gz = File::open(archive_file)?;
    let tar = GzDecoder::new(tar_gz);
    let mut archive = Archive::new(tar);

    log::info!(
        "Extracting archive {} to: {}",
        archive_file.display(),
        dest_dir.display()
    );
    archive.unpack(dest_dir)?;
    Ok(())
}

Very straightforward. However, this doesn't give us the same path that the Python version does. The python version returns the extracted path. So We're going to have to do that next.

fn extract_archive(archive_file: &PathBuf, dest_dir: &PathBuf) -> Result<PathBuf, Error> {
    let tar_gz = File::open(archive_file)?;
    let tar = GzDecoder::new(tar_gz);
    let mut archive = Archive::new(tar);

    log::info!(
        "Extracting archive {} to: {}",
        archive_file.display(),
        dest_dir.display()
    );
    let dir = {
        let entry = &archive
            .entries()?
            .next()
            .ok_or(Error::TarEntry("No entries in archive"))??;
        entry.path()?.into_owned()
    };
    let archive_dir = dest_dir.join(dir);
    if archive_dir.exists() {
        log::info!("Archive already extracted to: {}", archive_dir.display());
        return Ok(archive_dir);
    }

    let tar = archive.into_inner();
    let mut tar_gz = tar.into_inner();
    tar_gz.seek(io::SeekFrom::Start(0))?;
    let tar = GzDecoder::new(tar_gz);
    let mut archive = Archive::new(tar);
    archive.unpack(dest_dir)?;

    Ok(archive_dir)
}

This is a hack that I'm using in order to fetch the first entry in the tar archive, which is generally the top level directory stored inside. Then I have to unwind all the seeking I did by unwrapping the inner Reader, seeking to 0, then reconstructing the archive.

If anyone knows of a more sane way to do this, please let me know in the comments.

This also required me to create another variant for our Error enum, TarEntry.

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("reqwest error: {0}")]
    Reqwest(#[from] reqwest::Error),

    #[error("io error: {0}")]
    IO(#[from] std::io::Error),

    #[error("homedir error: {0}")]
    Home(#[from] homedir::GetHomeError),

    #[error("tar entry error: {0}")]
    TarEntry(&'static str),
}

I also threw in a condition to return early if the archive has already been extracted. In the future we may want to change this to use SHA-1 hashes to verify that the data is the same as what was downloaded.

Conclusion

Well, so far we've managed to download and extract our dataset to a centralized location. This is a good first step. The first line of our program looks very similar to that of the python version.

use std::path::PathBuf;

use color_eyre::eyre::{Context, Result};
use tardyai::{untar_images, Url};

fn main() -> Result<()> {
    env_logger::Builder::new()
        .filter_level(log::LevelFilter::Info)
        .init();
    color_eyre::install()?;

    let path: PathBuf = untar_images(Url::Pets)
        .context("downloading Pets")?
        .join("images");
    log::info!("Images are in: {}", path.display());
    Ok(())
}

In Part 3, we will figure out how to turn our images on disk into an ExactSizeDataset that can provide the images as Tensor structs, with their associated labels, and enable batching and other useful functions.

And if you want to see the code from this stage, you can either fetch the article-2 tag from git with

git co article-2

or browse it on github

Working through the fast.ai book in Rust - Part 1

Favil Orbedios — Sat, 18 Nov 2023 04:06:16 +0000

Introduction

I've been dreaming about completing the fast.ai course for about 4 years now. I've done some of the projects in the book, and it's fun, however, I've never liked the magical nature of the fastai library.

So I've decided to write this series as a means of understanding the underlying libraries and concepts. And since I love the language, and I'm a glutton for punishment, I've decided to do all of this in Rust. I've found the dfdx crate to be really approachable, and it has the same kind of strong typing that I love about the Rust language as a whole.

In particular, from the documentation:

dfdx is a cuda accelerated tensor and neural network library, writtten entirely in rust!

Additionally, it can track compile time shapes across tensor operations, ensuring that all your neural networks are checked at compile time.

I'm very excited about the compile time checking of the neural networks, but on the other hand, that also makes it more difficult to create neural networks on the fly. You have to make sure everything type checks. But once it does, you can be sure your dimensions line up properly.

By no means am I an expert on anything I'm talking about, this is as much of a learning experience for me as I assume it is for you, dear reader. That being said, if you want to read the ramblings of an engineer attempting to learn how to build Neural Networks in a language that is, by most metrics, terrible for the field, then please, by all means, read on.

Part 1 of this series will only talk about setting up the environment and showing off the Tensor math capabilities of the dfdx crate. We'll discuss actually setting up a neural network in part 2.

Setting up

As is customary with these types of articles, I must start with the obligatory explanation of how to install the tools I'm going to be using.

We'll just be using the standard Rust installation method. I'm running Linux so I'll copy the installation command below, Mac is the same. For Windows, you can click on the link in this paragraph and download the rustup-init.exe file appropriate for your CPU architecture.

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

This will install the rust compiler and the cargo toolchain that we will use to fetch dependencies, and build our code into native binaries.

Next, let's navigate to a comfortable directory where we can start coding. I use $HOME/git for all of my code.

cd $HOME/git/
mkdir -p articles/fastai-rust/
cd articles/fastai-rust/

Then we can use cargo to create our rust project. I'm just going to go by chapters for right now.

cargo new --bin chapter1
cd chapter1/

Now if you run cargo run in this folder you'll see the ceremonial "Hello, world!".

➜   cargo run
   Compiling chapter1 v0.1.0 (/home/klah/git/articles/fastai-rust/chapter1)
    Finished dev [unoptimized + debuginfo] target(s) in 0.25s
     Running `target/debug/chapter1`
Hello, world!

Let's do something slightly interesting in the next section.

Doing something interesting

Now "Hello, world!" is all well and good, but that's just a starting place. Let's figure out how to get something mildly interesting going. Right now, I'm thinking that doing some simple tensor math is a good starting place. We can get to the actual meat of the course after we've seen the type of thing that dfdx can do.

First let's add some useful crates, and our intended dfdx crate. I'm adding the env_logger and the log crates for ease of logging.

cargo add env_logger log dfdx

And if we edit our main.rs file a bit:

fn main() {
    env_logger::Builder::new()
        .filter_level(log::LevelFilter::Info)
        .init();
}

This is the magic incantation that I use a lot in order to get sane default logging. This is nice, because the default format includes a timestamp and the module that the log line was logged from. This can help with rudimentary timing at least to the second level granularity. And for a lot of what we're going to be doing, second level granularity will be nice, because it will take a long time to run some of the steps. And we won't have to inject specific timing code.

use dfdx::prelude::*;

fn main () {
    // ...

    let dev = AutoDevice::default();
}

This line creates a "device" that actually does the construction of tensors, and models. This can be either a CPU or an NVidia GPU running on Cuda. We need a device to create tensors, because if we are actually creating them on a GPU, we can't just create a vector in memory, we have to create it in the GPU memory. A device helps you create tensors and initialize them with data.

    // ... Rest of code
    let a: Tensor<Rank1<3>, f32, _> = dev.tensor([1.0, 2.0, 3.0]);
    let b: Tensor<Rank1<3>, f32, _> = dev.tensor([4.0, 5.0, 6.0]);

    let c = a.clone() + b.clone();

    log::info!("\n{:?} + {:?} = {:?}", a.array(), b.array(), c.array());

These lines create two one dimensional tensors and add them together, and log out the results. Tensors add together by just adding each individual element. So we can see from the output: [1.0, 2.0, 3.0] + [4.0, 5.0, 6.0] = [5.0, 7.0, 9.0]

Calling the .array() method turns the tensor into something that implements Debug in a sensible manor. In particular, it turns a one dimensional tensor into an array of the same size. We'll see how it similarily converts a multidimensional tensor into a multidimensional array in the next step.

    // ... Rest of code
    let e: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
    let f = dev.tensor([7.0, 8.0, 9.0]).broadcast();

    let g = e.clone().powf(0.5) + (f.clone() * 2.0);

    // Log the result
    log::info!("\n √({:?}) \n+ (2.0 * {:?}) \n= {:?}", e.array(), f.array(), g.array());

In these lines, I decided to do something different. The first line just creates a two dimensional tensor with a 2x3 shape. The line after also creates a two dimensional vector, but it does it by broadcasting a one dimensional tensor into the second dimension with the same shape as e.

This is done with fancy Rust type system trickery. Because we add them together and store them in g, Rust knows that they must have the same type, and therefore it knows to change the shape to a 2x3 tensor.

In the third line, I wanted to show some more advanced math operations. So I am taking the element-wise square root of e, and doubling f, and adding them together.

The final log line just makes the log easier to read by putting the tensors on their own line.

Putting it all together

So here I've collected all the lines together into a single program. You can also see all the code

use dfdx::prelude::*;

fn main() {
    // Initialize logging
    env_logger::Builder::new()
        .filter_level(log::LevelFilter::Info)
        .init();

    // Initialize dfdx device, either CPU or CUDA depending on availability
    let dev = AutoDevice::default();

    // Create two tensors and add them together
    let a: Tensor<Rank1<3>, f32, _> = dev.tensor([1.0, 2.0, 3.0]);
    let b: Tensor<Rank1<3>, f32, _> = dev.tensor([4.0, 5.0, 6.0]);

    let c = a.clone() + b.clone();

    // Log the result
    log::info!("\n{:?} + {:?} = {:?}", a.array(), b.array(), c.array());

    // It even supports higher dimensional tensors
    let e: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
    let f = dev.tensor([7.0, 8.0, 9.0]).broadcast();

    let g = e.clone().powf(0.5) + (f.clone() * 2.0);

    // Log the result
    log::info!("\n √({:?}) \n+ (2.0 * {:?}) \n= {:?}", e.array(), f.array(), g.array());
}

And if we run that we can see some simple tensor math getting performed.

➜   cargo run
   Compiling chapter1 v0.1.0 (/home/klah/git/articles/fastai-rust/chapter1)
    Finished dev [unoptimized + debuginfo] target(s) in 0.84s
     Running `target/debug/chapter1`
[2023-11-18T03:55:42Z INFO  chapter1]
    [1.0, 2.0, 3.0] + [4.0, 5.0, 6.0] = [5.0, 7.0, 9.0]
[2023-11-18T01:37:13Z INFO  chapter1]
     √([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    + (2.0 * [[7.0, 8.0, 9.0], [7.0, 8.0, 9.0]])
    = [[15.0, 17.414213, 19.73205], [16.0, 18.236069, 20.44949]]

Next time

In the next installment, I plan to start working on the dog/cat categorizer that is our first model from chapter 1.

This is not nearly as easy to do in Rust as it is in Python, so there will be a number of steps to get us into a position where we can actually create the model and run it.

Forem: Favil Orbedios

fast.ai Book in Rust - Chapter 2 - Part 2

Introduction

Saving weights is easy

Run times

Conclusion

fast.ai Book in Rust - Chapter 2 - Part 1

Introduction

Code Clean Up

Conclusion

Working through the fast.ai book in Rust - Part 6

Introduction

The Training Loop

Progress bars and stuff

Conclusion

Working through the fast.ai book in Rust - Part 5

Introduction

Some cleanup first

What about those weights?

What can we do with it?

Working through the fast.ai book in Rust - Part 4

Introduction

Let's build the model

What can we do about the weights

First source of frustration

Moving on to downloading models

Not so fast

Conclusion

Working through the fast.ai book in Rust - Part 3

Introduction

Refactoring

On to finding images

Let's get to the Tensors now

How do we optimize it?

Conclusion

Working through the fast.ai book in Rust - Part 2

Introduction

Creating a new Rust package

Let's add URLs

Actually download something why don't you?

Save it to the hard disk already

We have the archive, now what?

Conclusion

Working through the fast.ai book in Rust - Part 1

Introduction

Setting up

Doing something interesting

Putting it all together

Next time

Let's get to the `Tensors` now