apps/python/mlp/mlp.py - with_cpu - vuron/adept

apps/python/mlp/mlp.py
113 строк3 KB
kolkir
Add tests for Linear layer and update matmul usage
28 мар 2025, 23:53
28 мар 2025, 23:53bc788d8
import sys
from pathlib import Path
from tqdm import tqdm
import time

from adept import (
    Variable,
    Tensor,
    TensorProperties,
    Shape,
    device_t,
    dtype_t,
    no_grad,
)
from adept.nn import Linear, Module, relu
from adept.optim import SGD
from adept.data import MNISTDataset, CPPDataLoader
from adept.loss import cross_entropy_with_logits
from adept.serialize import FileInput, FileOutput

device = device_t.CPU
dtype = dtype_t.Float32

epochs = 5
batch_size = 32
lr = 0.01
lr_decay = 0.5


class MLP(Module):
    def __init__(self):
        super().__init__()
        self.l1 = Linear(28 * 28, 512, device, dtype)
        self.l2 = Linear(512, 256, device, dtype)
        self.l3 = Linear(256, 10, device, dtype)

    def forward(self, x):
        out = relu(self.l1(x))
        out = relu(self.l2(out))
        out = self.l3(out)
        return out


def main():
    if len(sys.argv) < 2:
        print("path to the MNIST dataset is missed!")
        exit(0)

    mnist_path = Path(sys.argv[1])
    if not mnist_path.exists():
        print("path to the MNIST dataset is incorrect!")
        exit(0)

    train_images_file = str(mnist_path / "train-images.idx3-ubyte")
    train_labels_file = str(mnist_path / "train-labels.idx1-ubyte")

    train_dataset = MNISTDataset(
        train_images_file, train_labels_file, dtype, device, True
    )
    train_dataloader = CPPDataLoader(train_dataset, batch_size)

    test_images_file = str(mnist_path / "t10k-images.idx3-ubyte")
    test_labels_file = str(mnist_path / "t10k-labels.idx1-ubyte")

    test_dataset = MNISTDataset(test_images_file, test_labels_file, dtype, device, True)
    test_dataloader = CPPDataLoader(test_dataset, batch_size)

    mlp = MLP()

    global lr
    if len(sys.argv) == 3 and Path(sys.argv[2]).exists():
        input = FileInput(sys.argv[2])
        mlp.load(input)
        lr = input.read("lr", dtype_t.Float32)

    optimizer = SGD(mlp.parameters(), lr)

    for epoch in tqdm(range(epochs), unit="epoch"):
        mlp.train()
        pbar = tqdm(train_dataloader, unit="batch")
        for b_i, batch in enumerate(pbar):
            x, y = batch
            out = mlp.forward(Variable(x, requires_grad=False))

            loss = cross_entropy_with_logits(out, Variable(y, requires_grad=False))
            if b_i % 64 == 0:
                pbar.set_postfix(loss=loss.data().float_at([0, 0]))

            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        optimizer.set_lr(optimizer.lr() * lr_decay)

        # save checkpoint
        checkpoint_path = f"checkpoint_{epoch}_{int(time.time()*1000.0)}.pt"
        output = FileOutput(checkpoint_path)
        mlp.save(output)
        output.write("lr", optimizer.lr())

        # test
        mlp.eval()
        with no_grad():
            total_loss = Tensor.zero(TensorProperties(Shape([1]), device, dtype))
            for b_i, batch in enumerate(train_dataloader):
                x, y = batch
                out = mlp.forward(Variable(x))
                loss = cross_entropy_with_logits(out, Variable(y))
                total_loss += loss.data()
            pbar.set_postfix(test_loss=total_loss.float_at([0]) / b_i)


if __name__ == "__main__":
    main()