Pytorch error: ValueError: pic should be 2/3 dimensional. Got 4 dimensions

Pytorch error: ValueError: pic should be 2/3 dimensional. Got 4 dimensions

正在尝试学习本教程 here。虽然当我 select 我的内容图像和样式图像时,当我尝试使用 imshow() 函数时我收到此错误:

ValueError: pic should be 2/3 dimensional. Got 4 dimensions.

使用 google 我还没有真正找到解决这个问题的方法。

这是我的代码:

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from PIL import Image
import matplotlib.pyplot as plt
import torchvision.transforms as transforms 
import torchvision.models as models
import copy
import numpy as np

# This detects if cuda is available for GPU training otherwise will use CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

# Desired size of the output image
imsize = 512 if torch.cuda.is_available() else 256
print(imsize)

# Helper function
def image_loader(image_name, imsize):
    # Scale the imported image and transform it into a torch tensor
    loader = transforms.Compose([transforms.Resize(imsize), transforms.ToTensor()])
    image = Image.open(image_name)
    # Fake batch dimension required to fit network's input dimension
    image = loader(image).unsqueeze(0)
    return image.to(device, torch.float)

# Helper function to show the tensor as a PIL image
def imshow(tensor, title=None):
    unloader = transforms.ToPILImage()
    image = tensor.cpu().clone()
    image = unloader(image)
    plt.imshow(image)
    if title is not None:
        plt.title(title)
    plt.pause(0.001) # Pause so that the plots are updated

# Loading of images
image_directory = './images/'
style_img = image_loader(image_directory + "pb.jpg", imsize)
content_img = image_loader(image_directory + "content.jpg", imsize)
assert style_img.size() == content_img.size(), "we need to import style and content images of the same size"

plt.figure()
imshow(style_img, title='style image')

任何建议都会很有帮助。

样式和内容图片供参考:

matplotlib.pyplot 需要 2D(灰度,尺寸=(W,H))或 3D(彩色,尺寸 =(W,H ,颜色通道)) 在 imshow-函数中。

您可能仍然将 batchsize 作为张量中的第一维,因为在您的代码中:

# Fake batch dimension required to fit network's input dimension
image = loader(image).unsqueeze(0)

它添加了第一个维度。如果是这样,请尝试使用:

plt.imshow(np.squeeze(image))

plt.imshow(image[0])