| import time |
|
|
| import mlx.core as mx |
| import mlx.nn |
| import mlx.optimizers as opt |
| import torch |
|
|
|
|
| def bench_mlx(steps: int = 20) -> float: |
| mx.set_default_device(mx.cpu) |
|
|
| class BenchNetMLX(mlx.nn.Module): |
| |
|
|
| def __init__(self, in_channels, hidden_channels=32): |
| super().__init__() |
|
|
| self.net = mlx.nn.Sequential( |
| mlx.nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1), |
| mlx.nn.ReLU(), |
| mlx.nn.Conv2d( |
| hidden_channels, 2 * hidden_channels, kernel_size=3, padding=1 |
| ), |
| mlx.nn.ReLU(), |
| mlx.nn.ConvTranspose2d( |
| 2 * hidden_channels, hidden_channels, kernel_size=3, padding=1 |
| ), |
| mlx.nn.ReLU(), |
| mlx.nn.ConvTranspose2d( |
| hidden_channels, in_channels, kernel_size=3, padding=1 |
| ), |
| ) |
|
|
| def __call__(self, input): |
| return self.net(input) |
|
|
| benchNet = BenchNetMLX(3) |
| mx.eval(benchNet.parameters()) |
| optim = opt.Adam(learning_rate=1e-3) |
|
|
| inputs = mx.random.normal([10, 256, 256, 3]) |
|
|
| params = benchNet.parameters() |
| optim.init(params) |
|
|
| state = [benchNet.state, optim.state] |
|
|
| def loss_fn(params, image): |
| benchNet.update(params) |
| pred_image = benchNet(image) |
| return (pred_image - image).abs().mean() |
|
|
| def step(params, image): |
| loss, grads = mx.value_and_grad(loss_fn)(params, image) |
| optim.update(benchNet, grads) |
| return loss |
|
|
| total_time = 0.0 |
| print("MLX:") |
| for i in range(steps): |
| start_time = time.perf_counter() |
|
|
| step(benchNet.parameters(), inputs) |
| mx.eval(state) |
| end_time = time.perf_counter() |
|
|
| print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms") |
| total_time += (end_time - start_time) * 1000 |
|
|
| return total_time |
|
|
|
|
| def bench_torch(steps: int = 20) -> float: |
| device = torch.device("cpu") |
|
|
| class BenchNetTorch(torch.nn.Module): |
| |
|
|
| def __init__(self, in_channels, hidden_channels=32): |
| super().__init__() |
|
|
| self.net = torch.nn.Sequential( |
| torch.nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1), |
| torch.nn.ReLU(), |
| torch.nn.Conv2d( |
| hidden_channels, 2 * hidden_channels, kernel_size=3, padding=1 |
| ), |
| torch.nn.ReLU(), |
| torch.nn.ConvTranspose2d( |
| 2 * hidden_channels, hidden_channels, kernel_size=3, padding=1 |
| ), |
| torch.nn.ReLU(), |
| torch.nn.ConvTranspose2d( |
| hidden_channels, in_channels, kernel_size=3, padding=1 |
| ), |
| ) |
|
|
| def forward(self, input): |
| return self.net(input) |
|
|
| benchNet = BenchNetTorch(3).to(device) |
| optim = torch.optim.Adam(benchNet.parameters(), lr=1e-3) |
|
|
| inputs = torch.randn(10, 3, 256, 256, device=device) |
|
|
| def loss_fn(pred_image, image): |
| return (pred_image - image).abs().mean() |
|
|
| total_time = 0.0 |
| print("PyTorch:") |
| for i in range(steps): |
| start_time = time.perf_counter() |
|
|
| optim.zero_grad() |
| pred_image = benchNet(inputs) |
| loss = loss_fn(pred_image, inputs) |
| loss.backward() |
| optim.step() |
|
|
| end_time = time.perf_counter() |
|
|
| print(f"{i:3d}, time={(end_time-start_time) * 1000:7.2f} ms") |
| total_time += (end_time - start_time) * 1000 |
|
|
| return total_time |
|
|
|
|
| def main(): |
| steps = 20 |
| time_mlx = bench_mlx(steps) |
| time_torch = bench_torch(steps) |
|
|
| print(f"average time of MLX: {time_mlx/steps:9.2f} ms") |
| print(f"total time of MLX: {time_mlx:9.2f} ms") |
| print(f"average time of PyTorch: {time_torch/steps:9.2f} ms") |
| print(f"total time of PyTorch: {time_torch:9.2f} ms") |
|
|
| diff = time_torch / time_mlx - 1.0 |
| print(f"torch/mlx diff: {100. * diff:+5.2f}%") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
