RuntimeError：张量a256的大小必须与非单维度3

如何解决RuntimeError：张量a256的大小必须与非单维度3

我刚刚训练了一个模型，现在对具有不同大小和类型（.tif 、、 bmp和.PNG）的自定义图像执行迭代。

我只是加载一个模型，想要生成三个图像，分别是压缩图像，原始图像和最终图像。当我放入训练文件时，以下代码可以正常工作，因为我在此处提供了CIFAR数据集进行测试。现在，我将提供自定义图像并制作单独的文件进行测试。

import torch
from model import End_to_end
from torch.autograd import Variable
from loss import loss_function
from grid import save_image
from torchvision import datasets,transforms

CUDA = torch.cuda.is_available()
if CUDA:
    model = End_to_end().cuda()
else:
    model = End_to_end()

EPOCHS = 20

testset = datasets.ImageFolder(root="/home/khawar/Desktop/End-to-End_IEEE-TVSCT/test/",transform=transforms.ToTensor())
model.load_state_dict(torch.load('./checkpoint/model.pth'))
print(testset.imgs)


def test(epoch):
    model.eval()
    test_loss = 0
    for i,(data,_) in enumerate(testset):
        data = Variable(data)
        final,residual_img,upscaled_image,com_img,orig_im = model(data.cuda())
        test_loss += loss_function(final,orig_im).data
        if epoch == EPOCHS and i == 0:
            #             save_image(final.data[0],'reconstruction_final',nrow=8)
            #             save_image(com_img.data[0],'com_img',nrow=8)
            n = min(data.size(0),6)
            print("saving the image " + str(n))
            comparison = torch.cat([data[:n],final[:n].cpu()])
            comparison = comparison.cpu()
            #             print(comparison.data)
            save_image(com_img[:n].data,'compressed_' + str(epoch) + '.png',nrow=n)
            save_image(comparison.data,'reconstruction_' + str(epoch) + '.png',nrow=n)

    test_loss /= len(testset.dataset)
    print('====> Test set loss: {:.4f}'.format(test_loss))


def save_images():
    epoch = EPOCHS
    model.eval()
    test_loss = 0
    for i,_) in enumerate(testset):
        data = Variable(data).unsqueeze(0)
        final,orig_im).data
        if i == 3:
            #             save_image(final.data[0],final[:n].cpu()])
            comparison = comparison.cpu()
            #             print(comparison.data)
            save_image(com_img[:1].data,'./compressed_image/compressed_' + str(i) + '.png',nrow=n)
            save_image(final[:1].data,'./final_image/final_' + str(epoch) + '.png',nrow=n)
            save_image(orig_im[:1].data,'./orginal_image/original_' + str(epoch) + '.png',nrow=n)

    test_loss /= len(testset.dataset)
    print('====> Test set loss: {:.4f}'.format(test_loss))


save_images()

Error
Traceback (most recent call last):
  File "/home/khawar/Desktop/End-to-End_IEEE-TVSCT/test.py",line 74,in <module>
    save_images()
  File "/home/khawar/Desktop/End-to-End_IEEE-TVSCT/test.py",line 51,in save_images
    for i,_) in enumerate(test_loader):
  File "/home/khawar/anaconda3/envs/End-to-End_IEEE-TVSCT/lib/python3.5/site-packages/torch/utils/data/dataloader.py",line 345,in __next__
    data = self._next_data()
  File "/home/khawar/anaconda3/envs/End-to-End_IEEE-TVSCT/lib/python3.5/site-packages/torch/utils/data/dataloader.py",line 385,in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/home/khawar/anaconda3/envs/End-to-End_IEEE-TVSCT/lib/python3.5/site-packages/torch/utils/data/_utils/fetch.py",line 47,in fetch
    return self.collate_fn(data)
  File "/home/khawar/anaconda3/envs/End-to-End_IEEE-TVSCT/lib/python3.5/site-packages/torch/utils/data/_utils/collate.py",line 79,in default_collate
    return [default_collate(samples) for samples in transposed]
  File "/home/khawar/anaconda3/envs/End-to-End_IEEE-TVSCT/lib/python3.5/site-packages/torch/utils/data/_utils/collate.py",in <listcomp>
    return [default_collate(samples) for samples in transposed]
  File "/home/khawar/anaconda3/envs/End-to-End_IEEE-TVSCT/lib/python3.5/site-packages/torch/utils/data/_utils/collate.py",line 55,in default_collate
    return torch.stack(batch,out=out)
RuntimeError: stack expects each tensor to be equal size,but got [3,288,352] at entry 0 and [3,256,256] at entry 1

解决方法

此问题是由于您具有不同形状的图像引起的。

使用torch.reshape将张量整形为固定尺寸（256、256、3）。

请参阅文档here