在jupyter notebook中如何对taichi计时

renzao_ai

于 2022-10-30 19:55:26 发布

阅读量172

点赞数

分类专栏：深度学习文章标签： python

本文链接：https://blog.csdn.net/zt5169/article/details/127603913

版权

深度学习专栏收录该内容

18 篇文章 1 订阅

订阅专栏

基本函数的计时

def is_prime(n: int):
    result = True
    for k in range(2,int(n**0.5)+1):
        if n%k == 0:
            result = False
            break
            
    return result

def count_primes(n: int) -> int:
    count = 0
    for k in range(2,n):
        if is_prime(k):
            count+=1
    return count

%time [print(count_primes(1000000))]

执行结果
78498
CPU times: total: 4.17 s
Wall time: 4.18 s

taichi on x86 cpu

import taichi as ti
ti.init()
@ti.func
def is_prime(n: int):
    result = True
    for k in range(2,int(n**0.5)+1):
        if n%k == 0:
            result = False
            break
            
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2,n):
        if is_prime(k):
            count+=1
    return count


%time [print(count_primes(100000000))]

执行结果
[Taichi] Starting on arch=x64
5761455
CPU times: total: 4min 5s
Wall time: 22.1 s

taichi on gpu rtx1060

import taichi as ti
ti.init(arch=ti.gpu)
@ti.func
def is_prime(n: int):
    result = True
    for k in range(2,int(n**0.5)+1):
        if n%k == 0:
            result = False
            break
            
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2,n):
        if is_prime(k):
            count+=1
    return count


%time [print(count_primes(10000000))]

执行结果
[Taichi] Starting on arch=cuda
664579
CPU times: total: 453 ms
Wall time: 452 ms