如何避免 GAN 生成手写数字代码阻塞主线程 (基于 PyTorch)

在使用 PyTorch 的 GAN 生成手写数字时，代码中的 plt.show() 会阻塞主线程，导致程序无法继续执行。为了解决这个问题，我们可以通过以下方法避免阻塞：

将 plt.show() 替换为 plt.savefig()： 将生成的图像保存到文件中，而不是直接显示，可以避免阻塞主线程。
使用子线程进行图像生成： 在主线程中调用 generate_and_save_image 函数，并将生成图像的过程放到子线程中进行。

以下代码示例展示了如何使用子线程解决 plt.show() 阻塞问题：

import threading
import matplotlib.pyplot as plt
import numpy as np

def generate_and_save_image(model, test_input):
    predictions = np.squeeze(model(test_input).cpu().numpy())
    fig = plt.figure(figsize=(4, 4))
    for i in range(predictions.shape[0]):
        plt.subplot(4, 4, i+1)
        plt.imshow((predictions[i]+1) / 2, cmap='gray')
        plt.axis('off')
    plt.savefig('generated_image.png')

def main():
    # create model and test_input
    t = threading.Thread(target=generate_and_save_image, args=(model, test_input))
    t.start()
    # continue main thread execution

if __name__ == '__main__':
    main()

通过以上方法，我们可以避免 plt.show() 阻塞主线程，并继续执行其他操作。