diff --git a/image_ref/main.py b/image_ref/main.py index 2b53852eb9671980350c6c4b77b5f8851734c017..1c3299b77c2e33b491d545cbc1ba25fb82d26710 100644 --- a/image_ref/main.py +++ b/image_ref/main.py @@ -214,7 +214,6 @@ def make_prediction_duo(model, data, f_name, f_name2): confidence_pred_list = [[] for i in range(n_class)] y_pred = [] y_true = [] - soft_max = nn.Softmax(dim=1) # iterate over test data for imaer, imana, img_ref, label in data: imaer = imaer.float() @@ -234,13 +233,12 @@ def make_prediction_duo(model, data, f_name, f_name2): imana = imana.cuda() img_ref = img_ref.cuda() label = label.cuda() - output = model(imaer, imana, img_ref) - confidence = soft_max(output) + confidence = model(imaer, imana, img_ref) print(label) print(confidence) confidence_pred_list[specie].append(confidence[:, 0].data.cpu().numpy()) # Mono class output (only most postive paire) - output = torch.argmax(output[:, 0]) + output = torch.argmax(confidence[:, 0]) label = torch.argmin(label) y_pred.append(output.tolist()) y_true.append(label.tolist()) # Save Truth diff --git a/image_ref/model.py b/image_ref/model.py index 5302e3b122d79d71372a781aee834f805840f5a4..0f702fb60cd5c8b96bf580b720b8d77a47f705c9 100644 --- a/image_ref/model.py +++ b/image_ref/model.py @@ -287,7 +287,7 @@ class Classification_model_duo_contrastive(nn.Module): self.im_encoder = resnet34(num_classes=2, in_channels=2) self.predictor = nn.Linear(in_features=2*2,out_features=2) - + self.soft_max = nn.Softmax(dim=1) def forward(self, input_aer, input_ana, input_ref): input_ana = torch.concat([input_ana, input_ref], dim=1) @@ -295,4 +295,5 @@ class Classification_model_duo_contrastive(nn.Module): out_aer = self.im_encoder(input_aer) out_ana = self.im_encoder(input_ana) out = torch.concat([out_aer,out_ana],dim=1) - return self.predictor(out) \ No newline at end of file + out = self.predictor(out) + return self.soft_max(out) \ No newline at end of file