souging/TRELLIS_TextTo3D

1import time2import torch3import vox2seq4 5 6if __name__ == "__main__":7    stats = {8        'z_order_cuda': [],9        'z_order_pytorch': [],10        'hilbert_cuda': [],11        'hilbert_pytorch': [],12    }13    RES = [16, 32, 64, 128, 256]14    for res in RES:15        coords = torch.meshgrid(torch.arange(res), torch.arange(res), torch.arange(res))16        coords = torch.stack(coords, dim=-1).reshape(-1, 3).int().cuda()17 18        start = time.time()19        for _ in range(100):20            code_z_cuda = vox2seq.encode(coords, mode='z_order').cuda()21        torch.cuda.synchronize()22        stats['z_order_cuda'].append((time.time() - start) / 100)23 24        start = time.time()25        for _ in range(100):26            code_z_pytorch = vox2seq.pytorch.encode(coords, mode='z_order').cuda()27        torch.cuda.synchronize()28        stats['z_order_pytorch'].append((time.time() - start) / 100)29 30        start = time.time()31        for _ in range(100):32            code_h_cuda = vox2seq.encode(coords, mode='hilbert').cuda()33        torch.cuda.synchronize()34        stats['hilbert_cuda'].append((time.time() - start) / 100)35 36        start = time.time()37        for _ in range(100):38            code_h_pytorch = vox2seq.pytorch.encode(coords, mode='hilbert').cuda()39        torch.cuda.synchronize()40        stats['hilbert_pytorch'].append((time.time() - start) / 100)41 42    print(f"{'Resolution':<12}{'Z-Order (CUDA)':<24}{'Z-Order (PyTorch)':<24}{'Hilbert (CUDA)':<24}{'Hilbert (PyTorch)':<24}")43    for res, z_order_cuda, z_order_pytorch, hilbert_cuda, hilbert_pytorch in zip(RES, stats['z_order_cuda'], stats['z_order_pytorch'], stats['hilbert_cuda'], stats['hilbert_pytorch']):44        print(f"{res:<12}{z_order_cuda:<24.6f}{z_order_pytorch:<24.6f}{hilbert_cuda:<24.6f}{hilbert_pytorch:<24.6f}")45 46