add eval writer

@@ -5,6 +5,7 @@ from pprint import pprint
 
 
 import torch
 import torch
 import torch.nn as nn
 import torch.nn as nn
+from torch.utils.tensorboard import SummaryWriter
 
 
 from utils import *
 from utils import *
 
 
@@ -35,7 +36,9 @@ def eval(model_path):
     test_loader = iter(get_dataloader(args, test_dataset))
     test_loader = iter(get_dataloader(args, test_dataset))
 
 
     print('\n[+] Start testing')
     print('\n[+] Start testing')
-    _test_res = validate(args, model, criterion, test_loader, step=0, writer=None)
+    log_dir = os.path.join('/content/drive/My Drive/CD2 Project/runs', model_name)
+    writer = SummaryWriter(log_dir=log_dir)
+    _test_res = validate(args, model, criterion, test_loader, step=0, writer=writer)
 
 
     print('\n[+] Valid results')
     print('\n[+] Valid results')
     print('  Acc@1 : {:.3f}%'.format(_test_res[0].data.cpu().numpy()[0]*100))
     print('  Acc@1 : {:.3f}%'.format(_test_res[0].data.cpu().numpy()[0]*100))
@@ -43,6 +46,7 @@ def eval(model_path):
     print('  Loss : {:.3f}'.format(_test_res[2].data))
     print('  Loss : {:.3f}'.format(_test_res[2].data))
     print('  Infer Time(per image) : {:.3f}ms'.format(_test_res[3]*1000 / len(test_dataset)))
     print('  Infer Time(per image) : {:.3f}ms'.format(_test_res[3]*1000 / len(test_dataset)))
 
 
+    writer.close()
 
 
 if __name__ == '__main__':
 if __name__ == '__main__':
     fire.Fire(eval)
     fire.Fire(eval)

@@ -54,6 +54,9 @@ def train(**kwargs):
     if torch.cuda.device_count() > 1:
     if torch.cuda.device_count() > 1:
         print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
         print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
         model = nn.DataParallel(model)
         model = nn.DataParallel(model)
+    
+    print('\n[+] Use {} GPUs'.format(torch.cuda.device_count()))
+    print('\n[+] Using GPU: {} '.format(torch.cuda.get_device_name(0)))
 
 
     start_t = time.time()
     start_t = time.time()
     for step in range(args.start_step, args.max_step):
     for step in range(args.start_step, args.max_step):
@@ -86,7 +89,7 @@ def train(**kwargs):
             print('  Acc@5 : {:.3f}%'.format(_valid_res[1].data.cpu().numpy()[0]*100))
             print('  Acc@5 : {:.3f}%'.format(_valid_res[1].data.cpu().numpy()[0]*100))
             print('  Loss : {}'.format(_valid_res[2].data))
             print('  Loss : {}'.format(_valid_res[2].data))
 
 
-            if _valid_res[0] > best_acc:
+            if _valid_res[0] >= best_acc:
                 best_acc = _valid_res[0]
                 best_acc = _valid_res[0]
                 torch.save(model.state_dict(), os.path.join(log_dir, "model","model.pt"))
                 torch.save(model.state_dict(), os.path.join(log_dir, "model","model.pt"))
                 print('\n[+] Model saved')
                 print('\n[+] Model saved')

@@ -86,7 +86,7 @@ def concat_image_features(image, features, max_features=3):
         feature = feature.view(1, h, w) #(3, h, w) input of size 3072
         feature = feature.view(1, h, w) #(3, h, w) input of size 3072
         # torch.Size([3, 32, 32])->[1, 32, 32]
         # torch.Size([3, 32, 32])->[1, 32, 32]
 
 
-        print("img_feature & feature size:\n", image_feature.size(),"\n", feature.size()) 
+        #print("img_feature & feature size:\n", image_feature.size(),"\n", feature.size()) 
         # img_feature & feature size:
         # img_feature & feature size:
         # torch.Size([1, 32, 32]) -> [1, 32, 64]
         # torch.Size([1, 32, 32]) -> [1, 32, 64]
         # torch.Size([3, 32, 32] ->[1, 32, 32]
         # torch.Size([3, 32, 32] ->[1, 32, 32]
@@ -395,7 +395,6 @@ def get_valid_transform(args, model):
 
 
 def train_step(args, model, optimizer, scheduler, criterion, batch, step, writer, device=None):
 def train_step(args, model, optimizer, scheduler, criterion, batch, step, writer, device=None):
     model.train()
     model.train()
-    #print('\nBatch\n', batch)
     images, target = batch
     images, target = batch
 
 
     if device:
     if device:
@@ -476,7 +475,7 @@ def validate(args, model, criterion, valid_loader, step, writer, device=None):
 
 
     return acc1, acc5, loss, infer_t
     return acc1, acc5, loss, infer_t
 
 
-
+#_acc1, _acc5 = accuracy(output, target, topk=(1, 5))
 def accuracy(output, target, topk=(1,)):
 def accuracy(output, target, topk=(1,)):
     """Computes the accuracy over the k top predictions for the specified values of k"""
     """Computes the accuracy over the k top predictions for the specified values of k"""
     with torch.no_grad():
     with torch.no_grad():