Remove log_softmax on inference

Author: synml

demo.py

@@ -2,7 +2,6 @@
 
 import matplotlib.colors
 import matplotlib.pyplot as plt
-import torch.nn.functional as F
 import torch.utils.data
 import tqdm
 
@@ -46,7 +45,6 @@
         with torch.cuda.amp.autocast(enabled=config['amp_enabled']):
             with torch.no_grad():
                 output = model(image)
-                output = F.log_softmax(output, dim=1)
                 output = torch.argmax(output, dim=1)
 
         # 1 배치단위 처리

eval.py

@@ -5,7 +5,6 @@
 import numpy as np
 import sklearn.metrics
 import torch.nn as nn
-import torch.nn.functional as F
 import torch.utils.data
 import tqdm
 
@@ -70,7 +69,6 @@ def evaluate(model, testloader, criterion, num_classes: int, amp_enabled: bool,
             val_loss += criterion(output, target).item()
 
             # Segmentation map 만들기
-            output = F.log_softmax(output, dim=1)
             output = torch.argmax(output, dim=1)
 
         # 혼동행렬 업데이트