fixed tests

src/test/resources/target_code/gluon/CNNSupervisedTrainer_mnist_mnistClassifier_net.py

@@ -352,21 +352,16 @@ class CNNSupervisedTrainer_mnist_mnistClassifier_net:
                             with open('src/test/resources/training_data/Show_attend_tell/dict.pkl', 'rb') as f:
                                 dict = pickle.load(f)
 
-
                         ax = fig.add_subplot(max_length//3, max_length//4, 1)
                         ax.imshow(train_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
 
                         for l in range(max_length):
                             attention = attentionList[l]
-                            attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1)
-                            attention = mx.nd.squeeze(attention)
+                            attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1).squeeze()
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
                             ax = fig.add_subplot(max_length//3, max_length//4, l+2)
                             if int(labels[l+1][0].asscalar()) > len(dict):
                                 ax.set_title("<unk>")
-                                img = ax.imshow(train_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
-                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
-                                break
                             elif dict[int(labels[l+1][0].asscalar())] == "<end>":
                                 ax.set_title(".")
                                 img = ax.imshow(train_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
@@ -374,14 +369,13 @@ class CNNSupervisedTrainer_mnist_mnistClassifier_net:
                                 break
                             else:
                                 ax.set_title(dict[int(labels[l+1][0].asscalar())])
-                                img = ax.imshow(train_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
-                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
-
+                            img = ax.imshow(train_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
+                            ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
 
                         plt.tight_layout()
                         target_dir = 'target/attention_images'
                         if not os.path.exists(target_dir):
-                                    os.makedirs(target_dir)
+                            os.makedirs(target_dir)
                         plt.savefig(target_dir + '/attention_train.png')
                         plt.close()
 
@@ -420,21 +414,16 @@ class CNNSupervisedTrainer_mnist_mnistClassifier_net:
                         fig = plt.figure(figsize=(15,15))
                         max_length = len(labels)-1
 
-
                         ax = fig.add_subplot(max_length//3, max_length//4, 1)
                         ax.imshow(test_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
 
                         for l in range(max_length):
                             attention = attentionList[l]
-                            attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1)
-                            attention = mx.nd.squeeze(attention)
+                            attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1).squeeze()
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
                             ax = fig.add_subplot(max_length//3, max_length//4, l+2)
                             if int(mx.nd.slice_axis(outputs[l+1], axis=0, begin=0, end=1).squeeze().asscalar()) > len(dict):
                                 ax.set_title("<unk>")
-                                img = ax.imshow(test_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
-                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
-                                break
                             elif dict[int(mx.nd.slice_axis(outputs[l+1], axis=0, begin=0, end=1).squeeze().asscalar())] == "<end>":
                                 ax.set_title(".")
                                 img = ax.imshow(test_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
@@ -442,9 +431,8 @@ class CNNSupervisedTrainer_mnist_mnistClassifier_net:
                                 break
                             else:
                                 ax.set_title(dict[int(mx.nd.slice_axis(outputs[l+1], axis=0, begin=0, end=1).squeeze().asscalar())])
-                                img = ax.imshow(test_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
-                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
-
+                            img = ax.imshow(test_images[0+test_batch_size*(batch_i)].transpose(1,2,0))
+                            ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
 
                         plt.tight_layout()
                         plt.savefig(target_dir + '/attention_test.png')