【Keras】 InvalidArgumentError: Incompatible shapesのエラーが出ます。

前提・実現したいこと

kerasを用いてVAEをCNNを用いて実装を試みました。
データはMNISTを使っています。
実行すると以下のようなエラーメッセージが表示されるのですが、どこが間違えているか特定できません。

発生している問題・エラーメッセージ

tensorflow.python.framework.errors_impl.InvalidArgumentError: Incompatible sh
apes: [10,28,28] vs. [10]

該当のソースコード

Python
1from keras.layers import Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Reshape, Flatten, Lambda
2from keras.models import Model
3from keras import backend as K
4import matplotlib.pyplot as plt
5import numpy as np
6
7intermediate_dim = 64
8latent_dim = 8
9
10class vae_util:
11    
12    def build_encoder(self):
13        input_img  = Input(shape=(28, 28, 1)) 
14        x = Conv2D(16, (3, 3), activation='relu', padding='same')(input_img)
15        x = MaxPooling2D((2, 2))(x)
16        x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
17        x = MaxPooling2D((2, 2))(x)
18        # (7, 7, 8)
19        x = Flatten()(x)
20        hidden = Dense(intermediate_dim, activation='relu')(x)
21        z_mean = Dense(latent_dim, activation='linear')(hidden)
22        z_log_var = Dense(latent_dim, activation='linear')(hidden)
23        return Model(input_img, [z_mean, z_log_var])
24    
25    def build_sampler(self):
26        z_mean = Input(shape=(latent_dim, ))
27        z_log_var = Input(shape=(latent_dim, ))
28        z = Lambda(self.sampling, output_shape=(latent_dim, ))([z_mean, z_log_var])
29        return Model([z_mean, z_log_var], z)
30
31    def build_decoder(self):
32        z = Input(shape=(latent_dim, ))
33        x = Dense(intermediate_dim, activation='relu')(z)
34        x = Dense(7*7*8, activation='relu')(x)
35        x = Reshape((7, 7, 8))(x)
36        x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
37        x = UpSampling2D((2, 2))(x)
38        x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
39        x = UpSampling2D((2, 2))(x)
40        decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)
41        return Model(z, decoded)
42    
43    def sampling(self, args):
44        z_mean, z_log_var = args
45        epsilon = np.random.normal(size=latent_dim)
46        return z_mean + K.sqrt(K.exp(z_log_var)) * epsilon
47    
48    def build_vae(self, encoder, sampler, decoder):
49        _, conv1, maxp1, conv2, maxp2, flat, dense1, encoder_mean, encoder_log_var = encoder.layers
50
51        input_img = Input(shape=(28, 28, 1))
52        x = conv1(input_img)
53        x = maxp1(x)
54        x = conv2(x)
55        x = maxp2(x)
56        x = flat(x)
57        x = dense1(x)
58        z_mean = encoder_mean(x)
59        z_log_var = encoder_log_var(x)
60
61        self.z_m = z_mean 
62        self.z_s = z_log_var 
63        
64        _, _, sampler_lambda = sampler.layers
65        z = sampler_lambda([z_mean, z_log_var])
66        
67        _, dense2, dense3, reshape1, conv3, upsampling1, conv4, upsampling2, conv5 = decoder.layers
68        
69        x = dense2(z)
70        x = dense3(x)
71        x = reshape1(x)
72        x = conv3(x)
73        x = upsampling1(x)
74        x = conv4(x)
75        x = upsampling2(x)
76        decoded = conv5(x)
77
78        return Model(input_img, decoded)
79    
80    def binary_crossentropy(self, y_true, y_pred):
81        return K.sum(K.binary_crossentropy(y_pred, y_true), axis=-1)
82
83
84    def vae_loss(self, x, x_decoded_mean):
85        z_mean = self.z_m
86        z_log_var = self.z_s
87
88        
89        # Reconstruction loss
90        rc_loss = K.sum(K.binary_crossentropy(x, x_decoded_mean), axis=-1)
91
92        # Regularization term (KL divergence)
93        kl_loss = -0.5 * K.sum(1 + z_log_var \
94                                 - K.square(z_mean) \
95                                 - K.exp(z_log_var), axis=-1)
96
97        # Average over mini-batch
98        return K.mean(rc_loss + kl_loss)
99
100        
101
102    def model_compile(self, model):
103        model.compile(optimizer='adam', loss=self.vae_loss)
104
105from keras.datasets import mnist
106import numpy as np
107
108(x_train, _), (x_test, _) = mnist.load_data()
109
110x_train = x_train.astype('float32') / 255.
111x_test = x_test.astype('float32') / 255.
112x_train = x_train[:, :, :, np.newaxis]
113x_test = x_test[:, :, :, np.newaxis]
114
115v = vae_util()
116encoder = v.build_encoder()
117sampler = v.build_sampler()
118decoder = v.build_decoder()
119vae= v.build_vae(encoder, sampler, decoder)
120v.model_compile(vae)
121
122epochs = 30
123result = vae.fit(x_train, x_train,
124                epochs=epochs,
125                batch_size=10,
126                shuffle=True,
127                validation_data=(x_test, x_test))
128
129
130