如何解决GAN生成手写数字代码阻塞主线程问题 (基于PyTorch)

在使用 PyTorch 训练 GAN 生成手写数字时，代码中可能会出现 plt.show() 导致主线程阻塞的问题。例如，以下代码用于生成并保存图像：

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.show()

可以使用 plt.draw() 替换 plt.show() 来解决阻塞问题。plt.draw() 只会将图像显示到缓存中，不会阻塞主线程。但是，需要手动关闭图像窗口才能继续执行程序。

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.draw()

需要注意的是，使用 plt.draw() 后，图像窗口不会自动关闭，需要手动关闭才能继续执行程序。