Как добавить слой изменения формы в nn.Sequential? ⇐ Python

[Anonymous]

Итак, я реализую генератор GAN, и мне нужна архитектура, показанная ниже:


Проблема в том, что когда я пытаюсь изменить форму вывода линейного слоя после BatchNorm и ReLU (на рис. Dense, поскольку они использовали Tensorflow), он выдает ошибку. ошибка как: TypeError: reshape(): аргумент «вход» (позиция 1) должен быть Tensor, а не int
Я понимаю ошибку, но не могу найти ее решение .
Есть ли другой способ изменить форму в nn.Sequential вместо явного вызова torch?
class Generator(nn.Module):

def __init__(self, z_dim=100, im_chan=1, hidden_dim=64, rdim=9216):

super(Generator, self).__init__()
self.z_dim = z_dim
self.gen = nn.Sequential(
nn.Linear(z_dim, rdim),
nn.BatchNorm2d(rdim,momentum=0.9),
nn.ReLU(inplace=True),
----> torch.reshape(rdim, (6,6,256)),
self.make_gen_block(rdim, hidden_dim*2),
self.make_gen_block(hidden_dim*2,hidden_dim),
self.make_gen_block(hidden_dim,im_chan,final_layer=True),
)
def make_gen_block(self, input_channels, output_channels, kernel_size=1, stride=2, final_layer=False):

if not final_layer:
return nn.Sequential(
nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
nn.BatchNorm2d(output_channels),
nn.ReLU(inplace=True)
)
else:
return nn.Sequential(
nn.ConvTranspose2d(input_channels, output_channels, kernel_size, stride),
nn.Tanh()
)

def unsqueeze_noise(self, noise):
return noise.view(len(noise), self.zdim, 1, 1)

def forward(self, noise):
x = self.unsqueeze_noise(noise)
return self.gen(x)

def get_noise(n_samples, z_dim, device='cpu'):
return torch.randn(n_samples, z_dim, device=device)

#Testing the Gen arch
gen = Generator()
num_test = 100

#test the hidden block
test_hidden_noise = get_noise(num_test, gen.z_dim)
test_hidden_block = gen.make_gen_block(6, 6, kernel_size=1,stride=2)
test_uns_noise = gen.unsqueeze_noise(test_hidden_noise)
hidden_output = test_hidden_block(test_uns_noise)


Подробнее здесь: https://stackoverflow.com/questions/741 ... sequential