Ben_Kosytorz
3a67c0979c
- Updated stablediffusion crate path from "../stable-diffusion-burn" to "./crates/stable-diffusion-burn" for proper workspace resolution - Enhanced .gitignore to include generated model files (.mpk, .pt, .bin, .safetensors, .ckpt) and user_data directory - Added Cargo.lock to gitignore with appropriate comment - Reorganized IDE files section in gitignore for better clarity - Added newline at end of file for proper formatting
216 lines
9.4 KiB
Rust
216 lines
9.4 KiB
Rust
use super::*;
|
|
use burn_tensor::{Bool, Tensor, TensorData};
|
|
|
|
#[test]
|
|
fn test_autodiff_checkpoint_complicated_computation() {
|
|
let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
|
|
let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
|
|
let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
|
|
let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
|
|
let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);
|
|
|
|
let device = Default::default();
|
|
let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
|
|
let tensor_1 = TestAutodiffTensor::from_data(data_1, &device).require_grad();
|
|
let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
|
|
let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
|
|
let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();
|
|
|
|
let tensor_5 = compute_bound_eager(tensor_0, tensor_1);
|
|
let tensor_6 = compute_bound_lazy(tensor_2, tensor_3.clone());
|
|
let tensor_7 = memory_bound_eager(tensor_3, tensor_4);
|
|
let tensor_8 = compute_bound_lazy(tensor_6, tensor_7.clone());
|
|
let tensor_9 = memory_bound_eager_scalar(tensor_7, 11.);
|
|
let tensor_10 = memory_bound_lazy(tensor_5, tensor_8.clone());
|
|
let tensor_11 = memory_bound_lazy(tensor_8, tensor_9);
|
|
let tensor_12 = compute_bound_lazy(tensor_10, tensor_11);
|
|
|
|
assert_checkpoint(tensor_12);
|
|
}
|
|
|
|
#[test]
|
|
fn test_autodiff_checkpoint_with_missing_requirement() {
|
|
let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
|
|
let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
|
|
|
|
let device = Default::default();
|
|
let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
|
|
let tensor_1 = TestAutodiffTensor::from_data(data_1, &device); // does not require_grad
|
|
|
|
let tensor_2 = memory_bound_eager(tensor_0, tensor_1);
|
|
let tensor_3 = memory_bound_eager_scalar(tensor_2.clone(), 11.);
|
|
let tensor_4 = memory_bound_eager_scalar(tensor_2.clone(), 11.);
|
|
let tensor_5 = compute_bound_lazy(tensor_3, tensor_4);
|
|
let tensor_6 = compute_bound_eager_scalar(tensor_5.clone(), 11.);
|
|
let tensor_7 = memory_bound_eager(tensor_5, tensor_2);
|
|
let tensor_8 = memory_bound_eager(tensor_6, tensor_7);
|
|
|
|
assert_checkpoint(tensor_8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_autodiff_checkpoint_with_many_duplicates() {
|
|
let data_0 = TensorData::from([[4.0, 7.0], [7.0, 7.0]]);
|
|
|
|
let device = Default::default();
|
|
let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
|
|
|
|
let tensor_1 = memory_bound_eager(tensor_0.clone(), tensor_0.clone());
|
|
let tensor_2 = compute_bound_eager(tensor_0.clone(), tensor_0.clone());
|
|
let tensor_3 = memory_bound_lazy(tensor_0.clone(), tensor_0.clone());
|
|
let tensor_4 = compute_bound_lazy(tensor_0.clone(), tensor_0.clone());
|
|
|
|
let tensor_5 = memory_bound_eager(tensor_1.clone(), tensor_0.clone());
|
|
let tensor_6 = memory_bound_eager(tensor_0.clone(), tensor_5.clone());
|
|
let tensor_7 = compute_bound_lazy(tensor_3.clone(), tensor_5.clone());
|
|
let tensor_8 = compute_bound_eager(tensor_4.clone(), tensor_2.clone());
|
|
let tensor_9 = memory_bound_lazy(tensor_6, tensor_7);
|
|
let tensor_10 = memory_bound_eager(tensor_0, tensor_9);
|
|
let tensor_11 = memory_bound_eager_scalar(tensor_10, 9.);
|
|
let tensor_12 = compute_bound_lazy(tensor_8, tensor_11);
|
|
|
|
assert_checkpoint(tensor_12);
|
|
}
|
|
|
|
#[test]
|
|
fn test_autodiff_checkpoint_with_long_chain_of_eager_memory_bound() {
|
|
let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
|
|
let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
|
|
let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
|
|
let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
|
|
let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);
|
|
|
|
let device = Default::default();
|
|
let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
|
|
let tensor_1 = TestAutodiffTensor::from_data(data_1, &device);
|
|
let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
|
|
let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
|
|
let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();
|
|
|
|
let tensor_5 = memory_bound_eager(tensor_0, tensor_1.clone());
|
|
let tensor_6 = memory_bound_eager(tensor_5, tensor_2);
|
|
let tensor_7 = memory_bound_eager(tensor_6, tensor_3);
|
|
let tensor_8 = memory_bound_eager(tensor_7, tensor_4);
|
|
let tensor_9 = memory_bound_lazy(tensor_8, tensor_1);
|
|
|
|
assert_checkpoint(tensor_9)
|
|
}
|
|
|
|
#[test]
|
|
fn test_autodiff_checkpoint_half_sub_graph_not_tracked() {
|
|
let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
|
|
let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
|
|
let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
|
|
let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
|
|
let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);
|
|
let data_5 = TensorData::from([[0.5, 7.0], [7.0, 7.0]]);
|
|
|
|
let device = Default::default();
|
|
let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device);
|
|
let tensor_1 = TestAutodiffTensor::from_data(data_1, &device);
|
|
let tensor_2 = TestAutodiffTensor::from_data(data_2, &device);
|
|
let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
|
|
let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();
|
|
let tensor_5 = TestAutodiffTensor::from_data(data_5, &device).require_grad();
|
|
|
|
let tensor_6 = memory_bound_lazy(tensor_0, tensor_1);
|
|
let tensor_7 = compute_bound_eager(tensor_6, tensor_2);
|
|
|
|
let tensor_8 = memory_bound_eager(tensor_3, tensor_4);
|
|
let tensor_9 = compute_bound_lazy(tensor_8, tensor_5);
|
|
|
|
let tensor_10 = compute_bound_lazy(tensor_7, tensor_9);
|
|
|
|
assert_checkpoint(tensor_10);
|
|
}
|
|
|
|
#[test]
|
|
fn test_autodiff_checkpoint_very_complex() {
|
|
let data_0 = TensorData::from([[0.0, 7.0], [7.0, 7.0]]);
|
|
let data_1 = TensorData::from([[0.1, 7.0], [7.0, 7.0]]);
|
|
let data_2 = TensorData::from([[0.2, 7.0], [7.0, 7.0]]);
|
|
let data_3 = TensorData::from([[0.3, 7.0], [7.0, 7.0]]);
|
|
let data_4 = TensorData::from([[0.4, 7.0], [7.0, 7.0]]);
|
|
|
|
let device = Default::default();
|
|
let tensor_0 = TestAutodiffTensor::<2>::from_data(data_0, &device).require_grad();
|
|
let tensor_1 = TestAutodiffTensor::from_data(data_1, &device);
|
|
let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
|
|
let tensor_3 = TestAutodiffTensor::from_data(data_3, &device).require_grad();
|
|
let tensor_4 = TestAutodiffTensor::from_data(data_4, &device).require_grad();
|
|
|
|
let tensor_5 = memory_bound_eager_scalar(tensor_0, 8.);
|
|
let tensor_6 = memory_bound_lazy(tensor_5.clone(), tensor_1.clone());
|
|
let tensor_7 = compute_bound_lazy(tensor_6.clone(), tensor_6);
|
|
let tensor_8 = memory_bound_lazy(tensor_1.clone(), tensor_5.clone());
|
|
let tensor_9 = memory_bound_eager_scalar(tensor_7.clone(), 7.);
|
|
let tensor_10 = compute_bound_eager(tensor_5, tensor_8);
|
|
let tensor_11 = memory_bound_eager(tensor_2.clone(), tensor_9);
|
|
let tensor_12 = memory_bound_lazy(tensor_2.clone(), tensor_2);
|
|
let tensor_13 = compute_bound_eager(tensor_10.clone(), tensor_11);
|
|
let tensor_14 = compute_bound_eager_scalar(tensor_3, 8.);
|
|
let tensor_15 = compute_bound_lazy(tensor_4, tensor_12);
|
|
let tensor_16 = memory_bound_lazy(tensor_10, tensor_7);
|
|
let tensor_17 = compute_bound_lazy(tensor_13, tensor_1);
|
|
let tensor_18 = memory_bound_eager(tensor_15, tensor_16);
|
|
let tensor_19 = compute_bound_eager(tensor_14, tensor_17);
|
|
let tensor_20 = memory_bound_lazy(tensor_18, tensor_19);
|
|
let tensor_21 = memory_bound_eager_scalar(tensor_20, 8.);
|
|
|
|
assert_checkpoint(tensor_21)
|
|
}
|
|
|
|
fn assert_checkpoint<const D: usize>(tensor: TestAutodiffTensor<D>) {
|
|
// Assert is not explicit here, but the test can fail
|
|
// - when a tensor is actually required more than n_required, it won't be found and will panic
|
|
// - when a tensor is actually required less than n_required, the backward states map won't be
|
|
// empty and will fail the assertion within the backward code, same for retro_forwards
|
|
tensor.backward();
|
|
}
|
|
|
|
// Does not save its state and does not need its parents
|
|
fn memory_bound_eager<const D: usize>(
|
|
tensor_a: TestAutodiffTensor<D>,
|
|
tensor_b: TestAutodiffTensor<D>,
|
|
) -> TestAutodiffTensor<D> {
|
|
tensor_a.add(tensor_b)
|
|
}
|
|
fn memory_bound_eager_scalar<const D: usize>(
|
|
tensor_a: TestAutodiffTensor<D>,
|
|
b: f32,
|
|
) -> TestAutodiffTensor<D> {
|
|
tensor_a.add_scalar(b)
|
|
}
|
|
|
|
// Saves its own state and does not need its parents
|
|
fn compute_bound_eager<const D: usize>(
|
|
tensor_a: TestAutodiffTensor<D>,
|
|
tensor_b: TestAutodiffTensor<D>,
|
|
) -> TestAutodiffTensor<D> {
|
|
let mask = Tensor::<TestAutodiffBackend, D, Bool>::empty(tensor_a.shape(), &tensor_a.device());
|
|
tensor_a.mask_where(mask, tensor_b)
|
|
}
|
|
fn compute_bound_eager_scalar<const D: usize>(
|
|
tensor_a: TestAutodiffTensor<D>,
|
|
b: f32,
|
|
) -> TestAutodiffTensor<D> {
|
|
let mask = Tensor::<TestAutodiffBackend, D, Bool>::empty(tensor_a.shape(), &tensor_a.device());
|
|
tensor_a.mask_fill(mask, b)
|
|
}
|
|
|
|
// Does not save its state and needs its parents
|
|
fn memory_bound_lazy<const D: usize>(
|
|
tensor_a: TestAutodiffTensor<D>,
|
|
tensor_b: TestAutodiffTensor<D>,
|
|
) -> TestAutodiffTensor<D> {
|
|
tensor_a.mul(tensor_b)
|
|
}
|
|
|
|
// Saves its own state and needs its parents
|
|
fn compute_bound_lazy<const D: usize>(
|
|
tensor_a: TestAutodiffTensor<D>,
|
|
tensor_b: TestAutodiffTensor<D>,
|
|
) -> TestAutodiffTensor<D> {
|
|
tensor_a.matmul(tensor_b)
|
|
}
|