fix : model contrastive

add : ref image .npy

config/config.py

@@ -32,7 +32,7 @@ def load_args_contrastive():
     parser.add_argument('--model', type=str, default='ResNet18')
     parser.add_argument('--model_type', type=str, default='duo')
     parser.add_argument('--dataset_dir', type=str, default='../data/processed_data/npy_image/data_training')
-    parser.add_argument('--dataset_ref_dir', type=str, default='image_ref/img_ref')
+    parser.add_argument('--dataset_ref_dir', type=str, default='img_ref')
     parser.add_argument('--output', type=str, default='../output/out_contrastive.csv')
     parser.add_argument('--save_path', type=str, default='../output/best_model_constrastive.pt')
     parser.add_argument('--pretrain_path', type=str, default=None)

dataset/dataset_ref.py

@@ -127,13 +127,15 @@ class ImageFolderDuo(data.Dataset):
         imgANA = self.loader(impathANA)
         label_ref = np.random.randint(0,len(self.classes)-1)
         class_ref = self.classes[label_ref]
-        path_ref = self.ref_dir + class_ref +'.npy'
+        path_ref = self.ref_dir +'/'+ class_ref + '.npy'
+        img_ref = self.loader(path_ref)
         if self.transform is not None:
             imgAER = self.transform(imgAER)
             imgANA = self.transform(imgANA)
+            img_ref = self.transform(img_ref)
         if self.target_transform is not None:
             target = 0 if self.target_transform(target) == label_ref else 1
-        img_ref = self.loader(path_ref)
+
         return imgAER, imgANA, img_ref, target
 
     def __len__(self):

image_ref/main.py

@@ -6,7 +6,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-from config.config import load_args
+from config.config import load_args, load_args_contrastive
 from dataset.dataset_ref import load_data, load_data_duo
 import torch
 import torch.nn as nn
@@ -146,13 +146,14 @@ def train_duo(model, data_train, optimizer, loss_function, epoch):
     for param in model.parameters():
         param.requires_grad = True
 
-    for imaer,imana, label in data_train:
+    for imaer,imana, img_ref, label in data_train:
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()
             imana = imana.cuda()
+            img_ref = img_ref.cuda()
             label = label.cuda()
-        pred_logits = model.forward(imaer,imana)
+        pred_logits = model.forward(imaer,imana,img_ref)
         pred_class = torch.argmax(pred_logits,dim=1)
         acc += (pred_class==label).sum().item()
         loss = loss_function(pred_logits,label)
@@ -172,13 +173,14 @@ def test_duo(model, data_test, loss_function, epoch):
     for param in model.parameters():
         param.requires_grad = False
 
-    for imaer,imana, label in data_test:
+    for imaer,imana, img_ref, label in data_test:
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()
             imana = imana.cuda()
+            img_ref = img_ref.cuda()
             label = label.cuda()
-        pred_logits = model.forward(imaer,imana)
+        pred_logits = model.forward(imaer,imana,img_ref)
         pred_class = torch.argmax(pred_logits,dim=1)
         acc += (pred_class==label).sum().item()
         loss = loss_function(pred_logits,label)
@@ -190,7 +192,7 @@ def test_duo(model, data_test, loss_function, epoch):
 
 def run_duo(args):
     #load data
-    data_train, data_test = load_data_duo(base_dir=args.dataset_dir, batch_size=args.batch_size, )
+    data_train, data_test = load_data_duo(base_dir=args.dataset_dir, batch_size=args.batch_size, ref_dir=args.dataset_ref_dir)
     #load model
     model = Classification_model_duo_contrastive(model = args.model, n_class=len(data_train.dataset.dataset.classes))
     model.double()
@@ -240,13 +242,14 @@ def make_prediction_duo(model, data, f_name):
     y_pred = []
     y_true = []
     # iterate over test data
-    for imaer,imana, label in data:
+    for imaer,imana,img_ref, label in data:
         label = label.long()
         if torch.cuda.is_available():
             imaer = imaer.cuda()
             imana = imana.cuda()
+            img_ref = img_ref.cuda()
             label = label.cuda()
-        output = model(imaer,imana)
+        output = model(imaer,imana,img_ref)
 
         output = (torch.max(torch.exp(output), 1)[1]).data.cpu().numpy()
         y_pred.extend(output)
@@ -277,7 +280,7 @@ def load_model(model, path):
 
 
 if __name__ == '__main__':
-    args = load_args()
+    args = load_args_contrastive()
     if args.model_type=='duo':
         run_duo(args)
     else :

image_ref/model.py

@@ -286,8 +286,8 @@ class Classification_model_duo_contrastive(nn.Module):
 
 
     def forward(self, input_aer, input_ana, input_ref):
-        input_ana = torch.concat(input_ana, input_ref, dim=2)
-        input_aer = torch.concat(input_aer, input_ref, dim=2)
+        input_ana = torch.concat([input_ana, input_ref], dim=1)
+        input_aer = torch.concat([input_aer, input_ref], dim=1)
         out_aer =  self.im_encoder(input_aer)
         out_ana = self.im_encoder(input_ana)
         out = torch.concat([out_aer,out_ana],dim=1)