如何并行化 Google Colab 笔记本以使其更快？

Question

我想知道是否有办法在 Google Colab 应用程序上并行化 Jupyter 笔记本。我使用 cobaya 包进行宇宙学分析，并且执行了许多蒙特卡洛马尔可夫链，因此我想知道如何并行化这些过程并使计算更快。

（我希望有一种方法可以将计算拆分到多个并行工作的核心中，即使是在 Google Colab 的虚拟环境中也是如此。是否需要付费帐户或免费版本就足够了？）

Answer 1

另一个更新：当因为一个所谓的糟糕答案而得到减分时，我感到非常难过，所以我想详细说明一下。

Python 并不是运行需要快速运行的密集计算的最佳语言。如果速度是最高优先级，请考虑 C/C++、Rust、Go 等编译语言，几乎所有语言都比 Python 快。

现在，如果你像我一样，你不喜欢这个答案，你宁愿找到一种加速 Python 的方法。

如果您还有大约 5 年的时间，并且确信新项目将兑现其承诺，您可以等待 mojo 编程语言 成熟并在快速环境中完成您的 Python 工作。

如果您不想等那么久，请使用现在可用的工具。 例如：使用 numba 库中的 njit 装饰器。 它将编译您的函数，以便它们可以运行得更快。 对于普通 Python 中的普通蒙特卡洛，这对我来说效果很好。马尔可夫链蒙特卡洛斯可能有一些额外的库，据我所知，这可能会给 njit 带来问题。

它的工作原理是这样的：

from numba import njit 

@njit
def my_super_fast_function():
    for i in range(50**50):
        1==1
    print("All finished")

这个警告可能非常烦人，以至于您想要深入研究并行处理或并行线程以使用机器的不同内核。我不是这方面的专家，所以请参阅其他 StackOverflow 问题的解释。

基本上，您为系统指定一个点来启动每个并行进程，并指定一个点来等待结果。然后您可以将结果放在一起，并希望速度有所提高...并不总是以可预测的方式工作...

更新：请参阅之前被问到的类似问题...以深入研究Python的本机多线程/多处理。

旧答案： 由于您似乎更多地寻找一般指南而不是特定解决方案，因此请检查指南（例如Quant Econ's GuideGoogle Colab 中的多处理笔记本）。

在上述指南中，他们使用 numba 模块及其 njit 装饰器来加速多线程......

示例取自指南：

from numba import njit, prange
import numpy as np

@njit # some function decorated with njit
def h(w, r=0.1, s=0.3, v1=0.1, v2=1.0):
    """
    Updates household wealth.
    """

    # Draw shocks
    R = np.exp(v1 * randn()) * (1 + r)
    y = np.exp(v2 * randn())

    # Update wealth
    w = R * s * w + y
    return w

@njit # long running function without pp
def compute_long_run_median(w0=1, T=1000, num_reps=50_000):

    obs = np.empty(num_reps)
    for i in range(num_reps):
        w = w0
        for t in range(T):
            w = h(w)
        obs[i] = w

    return np.median(obs)

@njit(parallel=True) # speedup by running with parallel flag
def compute_long_run_median_parallel(w0=1, T=1000, num_reps=50_000):

    obs = np.empty(num_reps)
    for i in prange(num_reps):
        w = w0
        for t in range(T):
            w = h(w)
        obs[i] = w

    return np.median(obs)