<input type="checkbox" id="" disabled=""

rust code for benching: <div class="highlight highlight-source-rust notranslate po

Python code for benching: <div class="highlight highlight-source-python notranslat

Both dfdx version and torch version should use flush to zero (<a class="issue-link js-

Plot benchmark speed against pytorch about dfdx HOT 5 CLOSED

coreylowman commented on May 13, 2024

Plot benchmark speed against pytorch

from dfdx.

Comments (5)

coreylowman commented on May 13, 2024

rust code for benching:

use dfdx::prelude::*;
use rand::{prelude::StdRng, SeedableRng};
use rand_distr::StandardNormal;
use std::time::{Duration, Instant};

fn main() {
    let mut rng = StdRng::seed_from_u64(0);

    let mut l: Linear<512, 256> = Default::default();
    l.randomize(&mut rng, &StandardNormal);

    let mut opt = Adam::default();

    const N: usize = 10000;
    let mut total = Duration::default();
    for _ in 0..N {
        let x: Tensor2D<32, 512> = Tensor2D::randn(&mut rng);
        let y = l.forward(x.traced());
        let loss = y.square().mean();
        let start = Instant::now();
        let gradients = loss.backward();
        opt.update(&mut l, gradients);
        total += start.elapsed();
    }
    println!("{:?} batch per s", N as f32 / total.as_secs_f32());
}

from dfdx.

coreylowman commented on May 13, 2024

Python code for benching:

from datetime import datetime, timedelta
import torch

torch.manual_seed(0)

l = torch.nn.Linear(512, 256)
opt = torch.optim.Adam(l.parameters())

total = timedelta()
N = 10000
for _ in range(N):
    x = torch.randn(32, 512)
    y = l(x)
    loss = y.square().mean()
    start = datetime.now()
    opt.zero_grad()
    loss.backward()
    opt.step()
    total += datetime.now() - start

print(N / total.total_seconds())