Unable to get all parameters of the model

I have created an autoencoder whose code is given below. But I try to check the parameters of the model, it does not show all the parameters of encoder and decoders.

class encoder(nn.Module):
	def __init__(self,
					input_dim,
					encoder_conv_filters,
					encoder_conv_kernel_size,
					encoder_conv_strides,
					z_dim,
					use_batch_norm = False,
					use_dropout = False):
		super(encoder, self).__init__()
		self.input_dim = input_dim
		self.encoder_conv_filters = encoder_conv_filters
		self.encoder_conv_kernel_size = encoder_conv_kernel_size
		self.encoder_conv_strides = encoder_conv_strides
		
		self.z_dim = z_dim
		self.use_batch_norm = use_batch_norm
		self.use_dropout = use_dropout
		
		self.n_layers_encoder = len(self.encoder_conv_filters)

		self.tensor_shapes = []
		self.conv2d = []
		self.batchNorm = []
		
		input_channels = self.input_dim[0]
		print(self.input_dim)
		for i in range(self.n_layers_encoder):
			'''if i > 0:
				padding = 1
			else:
				padding = (2, 1)'''
			padding = 1
			print(input_channels)
			self.conv2d.append(nn.Conv2d(in_channels = input_channels,
							out_channels = self.encoder_conv_filters[i],
							kernel_size = self.encoder_conv_kernel_size[i],
							stride = self.encoder_conv_strides[i],
							padding = padding))
			
			input_channels = self.encoder_conv_filters[i]

			self.batchNorm.append(nn.BatchNorm2d(self.encoder_conv_filters[i], affine = True))

	def forward(self, x):
		input_channels = x.shape[1]
		print("x shape:", x.shape)

		for i in range(self.n_layers_encoder):
			x = self.conv2d[i](x)
			x = nn.LeakyReLU()(x)

			if self.use_batch_norm:
				x = self.batchNorm[i](x)

			if self.use_dropout:
				x = nn.Dropout(p = 0.25)(x)
			
			print("x shape:", x.shape)
			self.tensor_shapes.append(x.size())
		
		self.shape_before_reshape = x.shape
		x = x.view(x.shape[0], -1)
		print("x shape after reshape:", x.shape)

		x = nn.Linear(x.shape[1], self.z_dim)(x)
		print("x shape:", x.shape)
		#x = nn.MaxPool2d(2, stride=2)(x)  # b, 16, 5, 5
		print('-'*100)
		
		return x, self.shape_before_reshape, self.tensor_shapes

class decoder(nn.Module):
	def __init__(self,
					decoder_conv_t_filters,
					decoder_conv_t_kernel_size,
					decoder_conv_t_strides,
					fitness_linear,
					z_dim,
					shape_before_reshape,
					tensor_shapes,
					use_batch_norm = False,
					use_dropout = False):
		super(decoder, self).__init__()
		
		self.decoder_conv_t_filters = decoder_conv_t_filters
		self.decoder_conv_t_kernel_size = decoder_conv_t_kernel_size
		self.decoder_conv_t_strides = decoder_conv_t_strides
		
		self.z_dim = z_dim

		self.fitness_linear = fitness_linear

		self.shape_before_reshape = shape_before_reshape
		self.tensor_shapes = tensor_shapes[::-1]
		self.tensor_shapes = self.tensor_shapes[1:]

		self.use_batch_norm = use_batch_norm
		self.use_dropout = use_dropout

		self.n_layers_decoder = len(self.decoder_conv_t_filters)

		self.convTranspose = []
		self.batchNorm = []

		self.linear1 = nn.Linear(self.z_dim, np.prod(self.shape_before_reshape[1:]))

		input_channels = shape_before_reshape[1]
		for i in range(self.n_layers_decoder):
			self.convTranspose.append(nn.ConvTranspose2d(in_channels = input_channels,
							out_channels = self.decoder_conv_t_filters[i],
							kernel_size = self.decoder_conv_t_kernel_size[i],
							stride = self.decoder_conv_t_strides[i],
							padding=1))
			input_channels = self.decoder_conv_t_filters[i]
			
			self.batchNorm.append(nn.BatchNorm2d(self.decoder_conv_t_filters[i], affine = True))

	def forward(self, x):
		print('='*100)
		print(self.tensor_shapes)
		print("\ndecoder: x shape:", x.shape)
		x = self.linear1(x)
		fitness = x
		
		print("decoder: x shape after nn.Linear:", x.shape)
		print("shape_before_reshape:", self.shape_before_reshape)
		x = x.view(-1, self.shape_before_reshape[1], self.shape_before_reshape[2], self.shape_before_reshape[3])
		print("decoder: x shape after reshape:", x.shape)
		
		input_channels = x.shape[1]

		for i in range(self.n_layers_decoder):
			if i == self.n_layers_decoder - 1:
				x = self.convTranspose[i](x) 
			else:
				x = self.convTranspose[i](x, output_size = self.tensor_shapes[i]) 
			
			if i < self.n_layers_decoder - 1:
				x = nn.LeakyReLU()(x)

				if self.use_batch_norm:
					x = self.batchNorm[i](x)

				if self.use_dropout:
					x = nn.Dropout(p = 0.25)(x)
			else:
				x = nn.Sigmoid()(x)
			
			print("i: {}, x shape:{}".format(i, x.shape))
		
		# Fitness estimation
		print("fitness:{}".format(fitness.shape))
		for output_shape in self.fitness_linear:
			fitness = nn.Linear(fitness.shape[1], output_shape)(fitness)
			fitness = nn.LeakyReLU()(fitness)
			
			if self.use_dropout:
				fitness = nn.Dropout(p = 0.25)(fitness)

			print("fitness:{}".format(fitness.shape))

		return x, fitness

class autoencoder(nn.Module):
	def __init__(self,
					input_dim,
					encoder_conv_filters,
					encoder_conv_kernel_size,
					encoder_conv_strides,
					decoder_conv_t_filters,
					decoder_conv_t_kernel_size,
					decoder_conv_t_strides,
					fitness_linear,
					z_dim,
					use_batch_norm = False,
					use_dropout = False):
		super(autoencoder, self).__init__()


		self.encoder = encoder(	input_dim,
										encoder_conv_filters,
										encoder_conv_kernel_size,
										encoder_conv_strides,
										z_dim,
										use_batch_norm,
										use_dropout)
		tmp_input = torch.randn(1, input_dim[0], input_dim[1], input_dim[2])
		_, shape_before_reshape, tensor_shapes = self.encoder(tmp_input)

		self.decoder = decoder(	decoder_conv_t_filters,
										decoder_conv_t_kernel_size,
										decoder_conv_t_strides,
										fitness_linear,
										z_dim,
										shape_before_reshape,
										tensor_shapes,
										use_batch_norm,
										use_dropout)

	def forward(self, x):
		x, shape_before_reshape, tensor_shapes  = self.encoder(x)
		x = self.decoder(x)
		return x

>>> print(model)
autoencoder(
  (encoder): encoder()
  (decoder): decoder(
    (linear1): Linear(in_features=2, out_features=512, bias=True)
  )
)
>>> print(model.state_dict().keys())
odict_keys(['decoder.linear1.weight', 'decoder.linear1.bias'])

To properly register modules in a list, you would have to use nn.ModuleList.
A plain Python list will not register the submodules inside the parent nn.Module and they will not be returned in e.g. model.parameters().

PS: I’ve formatted your code for easier debugging.
You can add code snippet by wrapping them into three backticks ``` :wink:

1 Like