rust-lstm 0.10.0

A complete LSTM neural network library with training capabilities, multiple optimizers, and peephole variants.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

"""PyTorch counterpart of `examples/benchmark.rs`.

    python validation/benchmark_pytorch.py --threads 1 --dtype float64
"""

import argparse
import statistics
import time

import torch

INPUT = 16
SEQ_LEN = 50


def median_micros(repeats, f):
    for _ in range(-(-repeats // 10)):
        f()
    samples = []
    for _ in range(repeats):
        start = time.perf_counter()
        f()
        samples.append((time.perf_counter() - start) * 1e6)
    return statistics.median(samples)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--threads", type=int, default=1)
    parser.add_argument("--dtype", choices=["float64", "float32"], default="float64")
    args = parser.parse_args()
    torch.set_num_threads(args.threads)
    dtype = getattr(torch, args.dtype)

    print("task,hidden,batch,microseconds")
    for hidden in (16, 64, 256):
        lstm = torch.nn.LSTM(INPUT, hidden, 1).to(dtype)
        x = torch.full((1, 1, INPUT), 0.1, dtype=dtype)
        state = (torch.zeros(1, 1, hidden, dtype=dtype),) * 2

        def step():
            with torch.no_grad():
                lstm(x, state)

        print(f"inference_step,{hidden},1,{median_micros(2000, step):.2f}")

        for batch in (1, 16):
            t = torch.arange(SEQ_LEN, dtype=dtype)[:, None, None] * 0.1
            xs = torch.sin(t).expand(SEQ_LEN, batch, INPUT).contiguous()
            ys = torch.cos(t).expand(SEQ_LEN, batch, hidden).contiguous()

            def train():
                lstm.zero_grad()
                out, _ = lstm(xs)
                ((out - ys) ** 2).mean(dim=(1, 2)).sum().backward()

            repeats = 30 if hidden == 256 else 200
            print(f"forward_backward_T{SEQ_LEN},{hidden},{batch},{median_micros(repeats, train):.2f}")


if __name__ == "__main__":
    main()