fix : dataloader batched

image_ref/main.py

@@ -62,8 +62,8 @@ def test_duo(model, data_test, loss_function, epoch):
         acc += (pred_class==label).sum().item()
         loss = loss_function(pred_logits,label)
         losses += loss.item()
-    losses = losses/len(data_test.dataset)
-    acc = acc/len(data_test.dataset)
+    losses = losses/(label.shape[0]*len(data_test.dataset))
+    acc = acc/(label.shape[0]*len(data_test.dataset))
     print('Test epoch {}, loss : {:.3f} acc : {:.3f}'.format(epoch,losses,acc))
     return losses,acc
 
@@ -112,23 +112,29 @@ def run_duo(args):
     plt.ylim(0, 1.05)
     plt.show()
 
-    plt.savefig('output/training_plot_contrastive_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
+    plt.savefig('../output/training_plot_contrastive_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
     #load and evaluate best model
     load_model(model, args.save_path)
-    make_prediction_duo(model,data_test, 'output/confusion_matrix_contractive_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
+    make_prediction_duo(model,data_test_batch, '../output/confusion_matrix_contractive_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model),
+                        '../output/confidence_matrix_contractive_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
 
 
-def make_prediction_duo(model, data, f_name):
+def make_prediction_duo(model, data, f_name, f_name2):
     for imaer, imana, img_ref, label in data:
-        n_class = len(label.shape[1])
+        n_class = label.shape[1]
         break
     confidence_pred_list = [[] for i in range(n_class)]
     y_pred = []
     y_true = []
     # iterate over test data
     for imaer,imana,img_ref, label in data:
+        imaer = imaer.transpose(0,1)
+        imana = imana.transpose(0,1)
+        img_ref = img_ref.transpose(0,1)
+        label = label.transpose(0,1)
+        label = label.squeeze()
         label = label.long()
-        specie = torch.argmax(label)
+        specie = torch.argmin(label)
 
         if torch.cuda.is_available():
             imaer = imaer.cuda()
@@ -136,7 +142,7 @@ def make_prediction_duo(model, data, f_name):
             img_ref = img_ref.cuda()
             label = label.cuda()
         output = model(imaer,imana,img_ref)
-        confidence_pred_list[specie].append(output.data.cpu().numpy())
+        confidence_pred_list[specie].append(output[:,0].data.cpu().numpy())
         output = (torch.max(torch.exp(output), 1)[1]).data.cpu().numpy()
         y_pred.extend(output)
 
@@ -145,6 +151,7 @@ def make_prediction_duo(model, data, f_name):
     # constant for classes
 
     # Build confusion matrix
+    classes = data.dataset.dataset.classes
     cf_matrix = confusion_matrix(y_true, y_pred)
     confidence_matrix = np.zeros((n_class,n_class))
     for i in range(n_class):
@@ -153,10 +160,19 @@ def make_prediction_duo(model, data, f_name):
     df_cm = pd.DataFrame(cf_matrix / np.sum(cf_matrix, axis=1)[:, None], index=[i for i in range(2)],
                          columns=['True','False'])
     print('Saving Confusion Matrix')
+    plt.clf()
     plt.figure(figsize=(14, 9))
     sn.heatmap(df_cm, annot=cf_matrix)
     plt.savefig(f_name)
 
+    df_cm = pd.DataFrame(confidence_matrix, index=[i for i in classes],
+                         columns=[i for i in classes])
+    print('Saving Confidence Matrix')
+    plt.clf()
+    plt.figure(figsize=(14, 9))
+    sn.heatmap(df_cm, annot=confidence_matrix)
+    plt.savefig(f_name2)
+
 
 def save_model(model, path):
     print('Model saved')