Module dfdx::tensor

source ·
Expand description

The Tensor struct, Cpu & Cuda devices, and traits like ZerosTensor, OnesTensor, SampleTensor.

At a high level a tensor is made up of:

  1. The crate::shapes::Shape of the array it stores
  2. The crate::shapes::Dtype of the elements of the array
  3. The Storage (e.g. Cpu or Cuda) that it uses to store the nd array
  4. A Tape, which can either actually be a tape (OwnedTape) or be empty (NoneTape).

Which are all generic parameters of Tensor. See the type’s docstring for more info

Creating a device

In order to do anything with tensors, you first need to construct the device that they will be stored on:

let dev: Cpu = Default::default();
let dev: Cpu = Cpu::seed_from_u64(0);
ⓘ
let dev: Cuda = Default::default();
let dev: Cuda = Cuda::seed_from_u64(1234);
let dev: Cuda = Cuda::try_build(0, 1234).unwrap();

Creating tensors

From arrays/vecs

See TensorFrom & TensorFromVec.

let _ = dev.tensor([1.0, 2.0, 3.0]);
let _ = dev.tensor_from_vec(vec![1.0, 2.0, 3.0], (3, ));

Filled with 0s or 1s

See ZerosTensor and OnesTensor.

let _: Tensor<Rank1<5>,f32 , _> = dev.zeros();
let _: Tensor<Rank2<3, 2>, f32, _> = dev.ones();

Filled with random data

See SampleTensor

let _: Tensor<Rank1<5>, f32, _> = dev.sample_uniform();
let _: Tensor<Rank2<3, 5>, f32, _> = dev.sample_normal();
// or pass in actual distributions
let _: Tensor<Rank1<3>, f32, _> = dev.sample(rand_distr::Standard);
let _: Tensor<Rank2<4, 3>, f32, _> = dev.sample(rand_distr::StandardNormal);

Copy data from slices

You can use Tensor::copy_from and Tensor::copy_into to copy data into a tensor:

let mut a: Tensor<Rank1<1000>, f32, _> = dev.zeros();
let buf: Vec<f32> = vec![1.0; 1000];
a.copy_from(&buf);

Modifying an already constructed tensor

There are only a few ways to do this, as normally you should just create a new tensor with tensor_ops.

See Tensor::fill_with_zeros, Tensor::fill_with_ones, Tensor::fill_with_distr

Converting tensors to rust arrays

Since the way tensors are stored is opaque to users (driven by whatever device the tensor is stored on), use the AsArray trait to convert tensors to actual rust arrays if you want to work with them directly.

let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
let t: [[f32; 3]; 2] = t.array();

Tracing gradients

Use the Tensor::trace or Tensor::traced methods to add OwnedTape to the Tensor. .trace() will clone the data, while .traced() will take ownership of the tensor and return a version with an OwnedTape.

Note that these two methods are only present for tensors without a tape already.

let t: Tensor<Rank1<5>,f32, _> = dev.zeros();
let mut grads = t.alloc_grads();
let t_clone: Tensor<Rank1<5>, f32, _, OwnedTape<f32, Cpu>> = t.trace(grads);

Gradient Accumulation

If you re-use the same gradients object without zero-ing out the gradients, you can implement gradient accumulation!

Serialization using numpy

See Tensor::save_to_npy and Tensor::load_from_npy.

You can also use Tensor::write_to_npz and Tensor::read_from_npz when working with zip archives.

Allocation Caching

By default, devices will cache allocations to reuse later. For example, the CPU will cache Vec<E> allocations in an hash map, and try to re-use the allocation when similarly sized data is used. This can result in significant speedups when memory allocation patterns are repetitive. If this results in extra memory use due to irregular allocation patterns there are two things you can do:

  1. Call Cache::empty_cache(), which will empty out all of the saved allocations.
  2. Disable the cache entirely by calling Cache::disable_cache(). This will empty out any existing allocations and prevent any new ones from being cached.

Modules

Structs

  • Cpu
    A device that stores data on the heap.
  • Cuda
    A Cuda device that enables constructing tensors on GPUs & running GPU kernels.
  • Gradients
    A generic container for keeping gradients of tensors keyed by the tensor’s UniqueId.
  • NoneTape
    Contains nothing. When Tape::add_backward_op is called, this struct does nothing.
  • OwnedTape
    Contains a Gradients and list of backward operations.
  • Tensor
    The single tensor struct that stores nd arrays and tapes.
  • UniqueId
    An id used in to associate gradients with Tensors.

Enums

Traits

Type Definitions