Added waitall to avoid crash when memory usage increases due to the asynchronous MXNet api

src/main/resources/templates/gluon/pythonExecuteTest.ftl

@@ -21,6 +21,8 @@
                     ${tc.getName(constant)} = mx.nd.full((test_batch_size, 1,), ${constant.intValue?c}, ctx=mx_context)
 </#list>
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
 <#list tc.architecture.networkInstructions as networkInstruction>

src/main/resources/templates/gluon/pythonExecuteTrain.ftl

@@ -21,6 +21,8 @@
                     ${tc.getName(constant)} = mx.nd.full((train_batch_size, 1,), ${constant.intValue?c}, ctx=mx_context)
 </#list>
 
+                    nd.waitall()
+
                     lossList = []
 
 <#list tc.architecture.networkInstructions as networkInstruction>

src/test/resources/target_code/CNNPredictor_Alexnet.h

@@ -44,7 +44,7 @@ public:
         MXPredGetOutputShape(handle, output_index, &shape, &shape_len);
         size = 1;
         for (mx_uint i = 0; i < shape_len; ++i) size *= shape[i];
-        assert(out_predictions_.size() == 1 || size == out_predictions_.size());
+        assert(size == out_predictions_.size());
         MXPredGetOutput(handle, output_index, &(out_predictions_[0]), out_predictions_.size());
 
     }

src/test/resources/target_code/CNNPredictor_CifarClassifierNetwork.h

@@ -44,7 +44,7 @@ public:
         MXPredGetOutputShape(handle, output_index, &shape, &shape_len);
         size = 1;
         for (mx_uint i = 0; i < shape_len; ++i) size *= shape[i];
-        assert(out_softmax_.size() == 1 || size == out_softmax_.size());
+        assert(size == out_softmax_.size());
         MXPredGetOutput(handle, output_index, &(out_softmax_[0]), out_softmax_.size());
 
     }

src/test/resources/target_code/CNNPredictor_VGG16.h

@@ -44,7 +44,7 @@ public:
         MXPredGetOutputShape(handle, output_index, &shape, &shape_len);
         size = 1;
         for (mx_uint i = 0; i < shape_len; ++i) size *= shape[i];
-        assert(out_predictions_.size() == 1 || size == out_predictions_.size());
+        assert(size == out_predictions_.size());
         MXPredGetOutput(handle, output_index, &(out_predictions_[0]), out_predictions_.size());
 
     }

src/test/resources/target_code/CNNSupervisedTrainer_Alexnet.py

@@ -289,6 +289,8 @@ class CNNSupervisedTrainer_Alexnet:
                     predictions_ = mx.nd.zeros((train_batch_size, 10,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     lossList = []
 
                     predictions_ = self._networks[0](data_)
@@ -322,7 +324,7 @@ class CNNSupervisedTrainer_Alexnet:
             train_test_iter.reset()
             metric = mx.metric.create(eval_metric, **eval_metric_params)
             for batch_i, batch in enumerate(train_test_iter):
-                if True: 
+                if True:
                     labels = [batch.label[i].as_in_context(mx_context) for i in range(1)]
 
                     data_ = batch.data[0].as_in_context(mx_context)
@@ -330,6 +332,8 @@ class CNNSupervisedTrainer_Alexnet:
                     predictions_ = mx.nd.zeros((test_batch_size, 10,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
                     predictions_ = self._networks[0](data_)
@@ -342,22 +346,32 @@ class CNNSupervisedTrainer_Alexnet:
                         logging.getLogger('matplotlib').setLevel(logging.ERROR)
 
                         plt.clf()
-                        fig = plt.figure(figsize=(10,10))
+                        fig = plt.figure(figsize=(15,15))
                         max_length = len(labels)-1
 
                         if(os.path.isfile('src/test/resources/training_data/Show_attend_tell/dict.pkl')):
                             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_resized = np.resize(attention.asnumpy(), (8, 8))
-                            ax = fig.add_subplot(max_length//3, max_length//4, l+1)
-                            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())
+                            ax = fig.add_subplot(max_length//3, max_length//4, l+2)
+                            if 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))
+                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
+                                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())
 
 
                         plt.tight_layout()
@@ -380,7 +394,7 @@ class CNNSupervisedTrainer_Alexnet:
             test_iter.reset()
             metric = mx.metric.create(eval_metric, **eval_metric_params)
             for batch_i, batch in enumerate(test_iter):
-                if True: 
+                if True:
                     labels = [batch.label[i].as_in_context(mx_context) for i in range(1)]
 
                     data_ = batch.data[0].as_in_context(mx_context)
@@ -388,6 +402,8 @@ class CNNSupervisedTrainer_Alexnet:
                     predictions_ = mx.nd.zeros((test_batch_size, 10,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
                     predictions_ = self._networks[0](data_)
@@ -397,19 +413,28 @@ class CNNSupervisedTrainer_Alexnet:
 
                     if save_attention_image == "True":
                         plt.clf()
-                        fig = plt.figure(figsize=(10,10))
+                        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=2, begin=0, end=1)
                             attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1)
                             attention = mx.nd.squeeze(attention)
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
-                            ax = fig.add_subplot(max_length//3, max_length//4, l+1)
-                            ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).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())
+                            ax = fig.add_subplot(max_length//3, max_length//4, l+2)
+                            if dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).asscalar())] == "<end>":
+                                ax.set_title(".")
+                                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
+                            else:
+                                ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).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())
 
 
                         plt.tight_layout()

src/test/resources/target_code/CNNSupervisedTrainer_CifarClassifierNetwork.py

@@ -289,6 +289,8 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
                     softmax_ = mx.nd.zeros((train_batch_size, 10,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     lossList = []
 
                     softmax_ = self._networks[0](data_)
@@ -322,7 +324,7 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
             train_test_iter.reset()
             metric = mx.metric.create(eval_metric, **eval_metric_params)
             for batch_i, batch in enumerate(train_test_iter):
-                if True: 
+                if True:
                     labels = [batch.label[i].as_in_context(mx_context) for i in range(1)]
 
                     data_ = batch.data[0].as_in_context(mx_context)
@@ -330,6 +332,8 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
                     softmax_ = mx.nd.zeros((test_batch_size, 10,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
                     softmax_ = self._networks[0](data_)
@@ -342,22 +346,32 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
                         logging.getLogger('matplotlib').setLevel(logging.ERROR)
 
                         plt.clf()
-                        fig = plt.figure(figsize=(10,10))
+                        fig = plt.figure(figsize=(15,15))
                         max_length = len(labels)-1
 
                         if(os.path.isfile('src/test/resources/training_data/Show_attend_tell/dict.pkl')):
                             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_resized = np.resize(attention.asnumpy(), (8, 8))
-                            ax = fig.add_subplot(max_length//3, max_length//4, l+1)
-                            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())
+                            ax = fig.add_subplot(max_length//3, max_length//4, l+2)
+                            if 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))
+                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
+                                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())
 
 
                         plt.tight_layout()
@@ -380,7 +394,7 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
             test_iter.reset()
             metric = mx.metric.create(eval_metric, **eval_metric_params)
             for batch_i, batch in enumerate(test_iter):
-                if True: 
+                if True:
                     labels = [batch.label[i].as_in_context(mx_context) for i in range(1)]
 
                     data_ = batch.data[0].as_in_context(mx_context)
@@ -388,6 +402,8 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
                     softmax_ = mx.nd.zeros((test_batch_size, 10,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
                     softmax_ = self._networks[0](data_)
@@ -397,19 +413,28 @@ class CNNSupervisedTrainer_CifarClassifierNetwork:
 
                     if save_attention_image == "True":
                         plt.clf()
-                        fig = plt.figure(figsize=(10,10))
+                        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=2, begin=0, end=1)
                             attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1)
                             attention = mx.nd.squeeze(attention)
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
-                            ax = fig.add_subplot(max_length//3, max_length//4, l+1)
-                            ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).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())
+                            ax = fig.add_subplot(max_length//3, max_length//4, l+2)
+                            if dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).asscalar())] == "<end>":
+                                ax.set_title(".")
+                                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
+                            else:
+                                ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).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())
 
 
                         plt.tight_layout()

src/test/resources/target_code/CNNSupervisedTrainer_VGG16.py

@@ -289,6 +289,8 @@ class CNNSupervisedTrainer_VGG16:
                     predictions_ = mx.nd.zeros((train_batch_size, 1000,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     lossList = []
 
                     predictions_ = self._networks[0](data_)
@@ -322,7 +324,7 @@ class CNNSupervisedTrainer_VGG16:
             train_test_iter.reset()
             metric = mx.metric.create(eval_metric, **eval_metric_params)
             for batch_i, batch in enumerate(train_test_iter):
-                if True: 
+                if True:
                     labels = [batch.label[i].as_in_context(mx_context) for i in range(1)]
 
                     data_ = batch.data[0].as_in_context(mx_context)
@@ -330,6 +332,8 @@ class CNNSupervisedTrainer_VGG16:
                     predictions_ = mx.nd.zeros((test_batch_size, 1000,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
                     predictions_ = self._networks[0](data_)
@@ -342,22 +346,32 @@ class CNNSupervisedTrainer_VGG16:
                         logging.getLogger('matplotlib').setLevel(logging.ERROR)
 
                         plt.clf()
-                        fig = plt.figure(figsize=(10,10))
+                        fig = plt.figure(figsize=(15,15))
                         max_length = len(labels)-1
 
                         if(os.path.isfile('src/test/resources/training_data/Show_attend_tell/dict.pkl')):
                             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_resized = np.resize(attention.asnumpy(), (8, 8))
-                            ax = fig.add_subplot(max_length//3, max_length//4, l+1)
-                            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())
+                            ax = fig.add_subplot(max_length//3, max_length//4, l+2)
+                            if 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))
+                                ax.imshow(attention_resized, cmap='gray', alpha=0.6, extent=img.get_extent())
+                                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())
 
 
                         plt.tight_layout()
@@ -380,7 +394,7 @@ class CNNSupervisedTrainer_VGG16:
             test_iter.reset()
             metric = mx.metric.create(eval_metric, **eval_metric_params)
             for batch_i, batch in enumerate(test_iter):
-                if True: 
+                if True:
                     labels = [batch.label[i].as_in_context(mx_context) for i in range(1)]
 
                     data_ = batch.data[0].as_in_context(mx_context)
@@ -388,6 +402,8 @@ class CNNSupervisedTrainer_VGG16:
                     predictions_ = mx.nd.zeros((test_batch_size, 1000,), ctx=mx_context)
 
 
+                    nd.waitall()
+
                     outputs = []
                     attentionList=[]
                     predictions_ = self._networks[0](data_)
@@ -397,19 +413,28 @@ class CNNSupervisedTrainer_VGG16:
 
                     if save_attention_image == "True":
                         plt.clf()
-                        fig = plt.figure(figsize=(10,10))
+                        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=2, begin=0, end=1)
                             attention = mx.nd.slice_axis(attention, axis=0, begin=0, end=1)
                             attention = mx.nd.squeeze(attention)
                             attention_resized = np.resize(attention.asnumpy(), (8, 8))
-                            ax = fig.add_subplot(max_length//3, max_length//4, l+1)
-                            ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).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())
+                            ax = fig.add_subplot(max_length//3, max_length//4, l+2)
+                            if dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).asscalar())] == "<end>":
+                                ax.set_title(".")
+                                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
+                            else:
+                                ax.set_title(dict[int(mx.nd.slice_axis(mx.nd.argmax(outputs[l+1], axis=1), axis=0, begin=0, end=1).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())
 
 
                         plt.tight_layout()