Hope for help, I have been dealing with th error for two weeks, but can not solve it. I am training my model with the input of images(3150150). The training model was finished, but I can not predict my test set, the code was got from zhaoyuzhi (Yuzhi ZHAO) · GitHub
Hierarchical-Regression-Network-for-Spectral-Reconstruction-from-RGB-Images
The test code was as follows
- test code
import argparse
import os
import torch
import numpy as np
import cv2
import hdf5storage as hdf5
import utils
import dataset
if name == “main”:
# ----------------------------------------
# Initialize the parameters
# ----------------------------------------
parser = argparse.ArgumentParser()
# Pre-train, saving, and loading parameters
parser.add_argument(’–pre_train’, type = bool, default = True, help = ‘pre_train or not’)
parser.add_argument(’–test_batch_size’, type = int, default = 1, help = ‘size of the testing batches for single GPU’)
parser.add_argument(’–num_workers’, type = int, default = 2, help = ‘number of cpu threads to use during batch generation’)
parser.add_argument(’–val_path’, type = str, default = ‘D:\NTIRE\HRNet\test’, help = ‘saving path that is a folder’)
parser.add_argument(’–task_name’, type = str, default = ‘track1’, help = ‘task name for loading networks, saving, and log’)
# Network initialization parameters
parser.add_argument(’–pad’, type = str, default = ‘reflect’, help = ‘pad type of networks’)
parser.add_argument(’–activ’, type = str, default = ‘lrelu’, help = ‘activation type of networks’)
parser.add_argument(’–norm’, type = str, default = ‘none’, help = ‘normalization type of networks’)
parser.add_argument(’–in_channels’, type = int, default = 3, help = ‘input channels for generator’)
parser.add_argument(’–out_channels’, type = int, default = 31, help = ‘output channels for generator’)
parser.add_argument(’–start_channels’, type = int, default = 64, help = ‘start channels for generator’)
parser.add_argument(’–init_type’, type = str, default = ‘xavier’, help = ‘initialization type of generator’)
parser.add_argument(’–init_gain’, type = float, default = 0.02, help = ‘initialization gain of generator’)
# Dataset parameters
parser.add_argument(’–baseroot’, type = str, default = ‘D:\NTIRE\HRNet\NTIRE2020_Test_Clean’, help = ‘baseroot’)
opt = parser.parse_args()
# ----------------------------------------
# Test
# ----------------------------------------
load_net_name_list = []
load_net_name_list.append('D:\\NTIRE\\HRNet\\track1\\G_epoch200_bs16.pth')
load_net_name_list.append('D:\\NTIRE\\HRNet\\track1\\G_epoch400_bs16.pth')
load_net_name_list.append('D:\\NTIRE\\HRNet\\track1\\G_epoch600_bs16.pth')
load_net_name_list.append('D:\\NTIRE\\HRNet\\track1\\G_epoch800_bs16.pth')
for i, name in enumerate(load_net_name_list):
if i < 7:
# Initialize
generator = utils.create_generator_val1(opt, name).cuda()
test_dataset = dataset.HS_multiscale_ValDSet(opt)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = opt.test_batch_size, shuffle = False, num_workers = opt.num_workers, pin_memory = True)
sample_folder = os.path.join(opt.val_path, opt.task_name, str(i))
utils.check_path(sample_folder)
print('Network name: %s; The %d-th iteration' % (name, i))
# forward
for j, (img1, img2, path) in enumerate(test_loader):
# To device
img1 = img1.cuda()
img2 = img2.cuda()
path = path[0]
print(opt.task_name, i, path)
# Forward propagation
with torch.no_grad():
out1 = generator(img1)
out2 = generator(img2)
out = torch.cat((out1, out2[:, :, 150:, :]), 2)
# Save
out = out.clone().data.permute(0, 2, 3, 1).cpu().numpy()[0, :, :, :].astype(np.float64)
save_img_name = path[:12] + '.mat'
save_img_path = os.path.join(sample_folder, save_img_name)
hdf5.write(data = out, path = 'cube', filename = save_img_path, matlab_compatible = True)
if i == 7:
# Initialize
generator = utils.create_generator_val2(opt, name).cuda()
test_dataset = dataset.HS_multiscale_ValDSet(opt)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = opt.test_batch_size, shuffle = False, num_workers = opt.num_workers, pin_memory = True)
sample_folder = os.path.join(opt.val_path, opt.task_name, str(i))
utils.check_path(sample_folder)
print('Network name: %s; The %d-th iteration' % (name, i))
# forward
for j, (img1, img2, path) in enumerate(test_loader):
# To device
img1 = img1.cuda()
img2 = img2.cuda()
path = path[0]
print(opt.task_name, i, path)
# Forward propagation
with torch.no_grad():
out1 = generator(img1)
out2 = generator(img2)
out = torch.cat((out1, out2[:, :, 150:, :]), 2)
# Save
out = out.clone().data.permute(0, 2, 3, 1).cpu().numpy()[0, :, :, :].astype(np.float64)
save_img_name = path[:12] + '.mat'
save_img_path = os.path.join(sample_folder, save_img_name)
hdf5.write(data = out, path = 'cube', filename = save_img_path, matlab_compatible = True)
The error is
- error
D:\python\python.exe D:/NTIRE/HRNet/test1.py
Generator is loaded!
Network name: D:\wangrenfa\model\NTIRE\HRNet\track1\G_epoch200_bs16.pth; The 0-th iteration
track1 0 ARAD_HS_0003_clean.png
Traceback (most recent call last):
File “D:/NTIRE/HRNet/test1.py”, line 70, in
out1 = generator(img1) # [0:480, 0:512, :], [1, 31, 480, 512]
File “D:\python\lib\site-packages\torch\nn\modules\module.py”, line 727, in _call_impl
result = self.forward(*input, **kwargs)
File “D:\NTIRE\HRNet\network_code1.py”, line 87, in forward
x2 = torch.cat((x2, x3), 1) # out: batch * (128 + 64) * 64 * 64
RuntimeError: Sizes of tensors must match except in dimension 2. Got 36 and 37 (The offending index is 0)
Process finished with exit code 1
The network_code1 is as follows:
network_code1
import torch
import torch.nn as nn
import torch.nn.functional as F
from network_module import *
import PixelUnShuffle
----------------------------------------
Initialize the networks
----------------------------------------
def weights_init(net, init_type = ‘normal’, init_gain = 0.02):
“”“Initialize network weights.
Parameters:
net (network) – network to be initialized
init_type (str) – the name of an initialization method: normal | xavier | kaiming | orthogonal
init_gain (float) – scaling factor for normal, xavier and orthogonal
In our paper, we choose the default setting: zero mean Gaussian distribution with a standard deviation of 0.02
“””
def init_func(m):
classname = m.class.name
if hasattr(m, ‘weight’) and classname.find(‘Conv’) != -1:
if init_type == ‘normal’:
torch.nn.init.normal_(m.weight.data, 0.0, init_gain)
elif init_type == ‘xavier’:
torch.nn.init.xavier_normal_(m.weight.data, gain = init_gain)
elif init_type == ‘kaiming’:
torch.nn.init.kaiming_normal_(m.weight.data, a = 0, mode = ‘fan_in’)
elif init_type == ‘orthogonal’:
torch.nn.init.orthogonal_(m.weight.data, gain = init_gain)
else:
raise NotImplementedError(‘initialization method [%s] is not implemented’ % init_type)
elif classname.find(‘BatchNorm2d’) != -1:
torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
torch.nn.init.constant_(m.bias.data, 0.0)
# apply the initialization function <init_func>
print('initialize network with %s type' % init_type)
net.apply(init_func)
----------------------------------------
Generator
----------------------------------------
class SGN(nn.Module):
def init(self, opt):
super(SGN, self).init()
# Top subnetwork, K = 3
self.top1 = Conv2dLayer(opt.in_channels * (4 ** 3), opt.start_channels * (2 ** 3), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.top21 = ResidualDenseBlock_5C(opt.start_channels * (2 ** 3), opt.start_channels * (2 ** 2), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.top22 = GlobalBlock(opt.start_channels * (2 ** 3), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm, sn = False, reduction = 4)
self.top3 = Conv2dLayer(opt.start_channels * (2 ** 3), opt.start_channels * (2 ** 3), 1, 1, 0, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
# Middle subnetwork, K = 2
self.mid1 = Conv2dLayer(opt.in_channels * (4 ** 2), opt.start_channels * (2 ** 2), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.mid2 = Conv2dLayer(int(opt.start_channels * (2 ** 2 + 2 ** 3 / 4)), opt.start_channels * (2 ** 2), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.mid31 = ResidualDenseBlock_5C(opt.start_channels * (2 ** 2), opt.start_channels * (2 ** 1), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.mid32 = GlobalBlock(opt.start_channels * (2 ** 2), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm, sn = False, reduction = 4)
self.mid4 = Conv2dLayer(opt.start_channels * (2 ** 2), opt.start_channels * (2 ** 2), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
# Bottom subnetwork, K = 1
self.bot1 = Conv2dLayer(opt.in_channels * (4 ** 1), opt.start_channels * (2 ** 1), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.bot2 = Conv2dLayer(int(opt.start_channels * (2 ** 1 + 2 ** 2 / 4)), opt.start_channels * (2 ** 1), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.bot31 = ResidualDenseBlock_5C(opt.start_channels * (2 ** 1), opt.start_channels * (2 ** 0), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.bot32 = ResidualDenseBlock_5C(opt.start_channels * (2 ** 1), opt.start_channels * (2 ** 0), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.bot33 = GlobalBlock(opt.start_channels * (2 ** 1), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm, sn = False, reduction = 4)
self.bot4 = Conv2dLayer(opt.start_channels * (2 ** 1), opt.start_channels * (2 ** 1), 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
# Mainstream
self.main1 = Conv2dLayer(opt.in_channels, opt.start_channels, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.main2 = Conv2dLayer(int(opt.start_channels * (2 ** 0 + 2 ** 1 / 4)), opt.start_channels, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.main31 = ResidualDenseBlock_5C(opt.start_channels, opt.start_channels // 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.main32 = ResidualDenseBlock_5C(opt.start_channels, opt.start_channels // 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.main33 = ResidualDenseBlock_5C(opt.start_channels, opt.start_channels // 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.main34 = ResidualDenseBlock_5C(opt.start_channels, opt.start_channels // 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.main35 = GlobalBlock(opt.start_channels, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm, sn = False, reduction = 4)
self.main4 = Conv2dLayer(opt.start_channels, opt.out_channels, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
def forward(self, x):
# PixelUnShuffle input: batch * 3 * 256 * 256
x1 = PixelUnShuffle.pixel_unshuffle(x, 2) # out: batch * 12 * 128 * 128
x2 = PixelUnShuffle.pixel_unshuffle(x, 4) # out: batch * 48 * 64 * 64
x3 = PixelUnShuffle.pixel_unshuffle(x, 8) # out: batch * 192 * 32 * 32
# Top subnetwork suppose the start_channels = 32
x3 = self.top1(x3) # out: batch * 256 * 32 * 32
x3 = self.top21(x3) # out: batch * 256 * 32 * 32
x3 = self.top22(x3) # out: batch * 256 * 32 * 32
x3 = self.top3(x3) # out: batch * 256 * 32 * 32
x3 = F.pixel_shuffle(x3, 2) # out: batch * 64 * 64 * 64, ready to be concatenated
# Middle subnetwork
x2 = self.mid1(x2) # out: batch * 128 * 64 * 64
x2 = torch.cat((x2, x3), 1) # out: batch * (128 + 64) * 64 * 64
x2 = self.mid2(x2) # out: batch * 128 * 64 * 64
x2 = self.mid31(x2) # out: batch * 128 * 64 * 64
x2 = self.mid32(x2) # out: batch * 128 * 64 * 64
x2 = self.mid4(x2) # out: batch * 128 * 64 * 64
x2 = F.pixel_shuffle(x2, 2) # out: batch * 32 * 128 * 128, ready to be concatenated
# Bottom subnetwork
x1 = self.bot1(x1) # out: batch * 64 * 128 * 128
x1 = torch.cat((x1, x2), 1) # out: batch * (64 + 32) * 128 * 128
x1 = self.bot2(x1) # out: batch * 64 * 128 * 128
x1 = self.bot31(x1) # out: batch * 64 * 128 * 128
x1 = self.bot32(x1) # out: batch * 64 * 128 * 128
x1 = self.bot33(x1) # out: batch * 64 * 128 * 128
x1 = self.bot4(x1) # out: batch * 64 * 128 * 128
x1 = F.pixel_shuffle(x1, 2) # out: batch * 16 * 256 * 256, ready to be concatenated
# U-Net generator with skip connections from encoder to decoder
x = self.main1(x) # out: batch * 32 * 256 * 256
x = torch.cat((x, x1), 1) # out: batch * (32 + 16) * 256 * 256
x = self.main2(x) # out: batch * 32 * 256 * 256
x = self.main31(x) # out: batch * 32 * 256 * 256
x = self.main32(x) # out: batch * 32 * 256 * 256
x = self.main33(x) # out: batch * 32 * 256 * 256
x = self.main34(x) # out: batch * 32 * 256 * 256
x = self.main35(x) # out: batch * 32 * 256 * 256
x = self.main4(x) # out: batch * 3 * 256 * 256
return x
‘’’
class SubNet(nn.Module):
def init(self, opt):
super(SubNet, self).init()
self.E1 = Conv2dLayer(opt.in_channels, opt.start_channels, 7, 1, 3, pad_type = opt.pad, activation = opt.activ, norm = ‘none’)
self.E2 = Conv2dLayer(opt.start_channels, opt.start_channels * 2, 4, 2, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.E3 = Conv2dLayer(opt.start_channels * 2, opt.start_channels * 4, 4, 2, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB1 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB2 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB3 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB4 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB5 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB6 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB7 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB8 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.D1 = TransposeConv2dLayer(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.D2 = TransposeConv2dLayer(opt.start_channels * 2, opt.start_channels, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.TM = Conv2dLayer(opt.start_channels, opt.start_channels, 1, 1, 0, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.D3 = Conv2dLayer(opt.start_channels, opt.out_channels, 7, 1, 3, pad_type = opt.pad, activation = ‘sigmoid’, norm = ‘none’)
def forward(self, x):
# Encoder
x = self.E1(x) # out: batch * 64 * 128 * 128
x = self.E2(x) # out: batch * 128 * 64 * 64
x = self.E3(x) # out: batch * 256 * 32 * 32
# Bottleneck
x = self.RDB1(x) # out: batch * 256 * 32 * 32
x = self.RDB2(x) # out: batch * 256 * 32 * 32
x = self.RDB3(x) # out: batch * 256 * 32 * 32
x = self.RDB4(x) # out: batch * 256 * 32 * 32
x = self.RDB5(x) # out: batch * 256 * 32 * 32
x = self.RDB6(x) # out: batch * 256 * 32 * 32
x = self.RDB7(x) # out: batch * 256 * 32 * 32
x = self.RDB8(x) # out: batch * 256 * 32 * 32
# Decoder
x = self.D1(x) # out: batch * 128 * 64 * 64
x = self.D2(x) # out: batch * 64 * 128 * 128
tm = self.TM(x) # out: batch * 64 * 128 * 128
x = self.D3(tm) # out: batch * 3 * 128 * 128
return tm, x
‘’’
‘’’
class SubNet(nn.Module):
def init(self, opt):
super(SubNet, self).init()
self.E1 = Conv2dLayer(opt.in_channels, opt.start_channels, 7, 1, 3, pad_type = opt.pad, activation = opt.activ, norm = ‘none’)
self.E2 = Conv2dLayer(opt.start_channels, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.E3 = Conv2dLayer(opt.start_channels * 2, opt.start_channels * 4, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB1 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB2 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB3 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB4 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB5 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB6 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB7 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.RDB8 = ResidualDenseBlock_5C(opt.start_channels * 4, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.D1 = Conv2dLayer(opt.start_channels, opt.start_channels * 2, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.D2 = Conv2dLayer(opt.start_channels * 2, opt.start_channels, 3, 1, 1, pad_type = opt.pad, activation = opt.activ, norm = opt.norm)
self.D3 = Conv2dLayer(opt.start_channels, opt.out_channels, 7, 1, 3, pad_type = opt.pad, activation = ‘sigmoid’, norm = ‘none’)
def forward(self, x):
# Encoder
x = self.E1(x) # out: batch * 64 * 128 * 128
x = self.E2(x) # out: batch * 128 * 128 * 128
x = self.E3(x) # out: batch * 256 * 128 * 128
# Bottleneck
x = self.RDB1(x) # out: batch * 256 * 128 * 128
x = self.RDB2(x) # out: batch * 256 * 128 * 128
x = self.RDB3(x) # out: batch * 256 * 128 * 128
x = self.RDB4(x) # out: batch * 256 * 128 * 128
x = self.RDB5(x) # out: batch * 256 * 128 * 128
x = self.RDB6(x) # out: batch * 256 * 128 * 128
x = self.RDB7(x) # out: batch * 256 * 128 * 128
x = self.RDB8(x) # out: batch * 256 * 128 * 128
# Decoder
x = self.D1(x) # out: batch * 128 * 128 * 128
tm = self.D2(x) # out: batch * 64 * 128 * 128
x = self.D3(tm) # out: batch * 3 * 128 * 128
return tm, x
‘’’