EmbedNet.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. #!/usr/bin/python
  2. #-*- coding: utf-8 -*-
  3. import torch, pickle
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import numpy, math, pdb, sys
  7. import time, importlib
  8. from DatasetLoader import test_dataset_loader
  9. from torch.cuda.amp import autocast, GradScaler
  10. import cv2
  11. import glob
  12. from PIL import Image
  13. from flask import Flask, request
  14. class EmbedNet(nn.Module):
  15. def __init__(self, model, optimizer, trainfunc, nPerClass, **kwargs):
  16. super(EmbedNet, self).__init__();
  17. ## __S__ is the embedding model
  18. EmbedNetModel = importlib.import_module('models.'+model).__getattribute__('MainModel')
  19. self.__S__ = EmbedNetModel(**kwargs);
  20. ## __L__ is the classifier plus the loss function
  21. LossFunction = importlib.import_module('loss.'+trainfunc).__getattribute__('LossFunction')
  22. self.__L__ = LossFunction(**kwargs);
  23. ## Number of examples per identity per batch
  24. self.nPerClass = nPerClass
  25. def forward(self, data, label=None):
  26. data = data.reshape(-1,data.size()[-3],data.size()[-2],data.size()[-1])
  27. outp = self.__S__.forward(data)
  28. if label == None:
  29. return outp
  30. else:
  31. outp = outp.reshape(self.nPerClass,-1,outp.size()[-1]).transpose(1,0).squeeze(1)
  32. nloss, prec1 = self.__L__.forward(outp,label)
  33. return nloss, prec1
  34. class ModelTrainer(object):
  35. def __init__(self, embed_model, optimizer, scheduler, mixedprec, **kwargs):
  36. self.__model__ = embed_model
  37. ## Optimizer (e.g. Adam or SGD)
  38. Optimizer = importlib.import_module('optimizer.'+optimizer).__getattribute__('Optimizer')
  39. self.__optimizer__ = Optimizer(self.__model__.parameters(), **kwargs)
  40. ## Learning rate scheduler
  41. Scheduler = importlib.import_module('scheduler.'+scheduler).__getattribute__('Scheduler')
  42. self.__scheduler__, self.lr_step = Scheduler(self.__optimizer__, **kwargs)
  43. ## For mixed precision training
  44. self.scaler = GradScaler()
  45. self.mixedprec = mixedprec
  46. assert self.lr_step in ['epoch', 'iteration']
  47. # ## ===== ===== ===== ===== ===== ===== ===== =====
  48. # ## Train network
  49. # ## ===== ===== ===== ===== ===== ===== ===== =====
  50. def train_network(self, loader, verbose):
  51. self.__model__.train();
  52. stepsize = loader.batch_size;
  53. counter = 0;
  54. index = 0;
  55. loss = 0;
  56. top1 = 0 # EER or accuracy
  57. tstart = time.time()
  58. for data, label in loader:
  59. data = data.transpose(1,0)
  60. ## Reset gradients
  61. self.__model__.zero_grad();
  62. ## Forward and backward passes
  63. if self.mixedprec:
  64. with autocast():
  65. nloss, prec1 = self.__model__(data.cuda(), label.cuda())
  66. self.scaler.scale(nloss).backward();
  67. self.scaler.step(self.__optimizer__);
  68. self.scaler.update();
  69. else:
  70. nloss, prec1 = self.__model__(data.cuda(), label.cuda())
  71. nloss.backward();
  72. self.__optimizer__.step();
  73. loss += nloss.detach().cpu();
  74. top1 += prec1.detach().cpu();
  75. counter += 1;
  76. index += stepsize;
  77. telapsed = time.time() - tstart
  78. tstart = time.time()
  79. if verbose:
  80. sys.stdout.write("\rProcessing (%d) "%(index));
  81. sys.stdout.write("Loss %f TEER/TAcc %2.3f%% - %.2f Hz "%(loss/counter, top1/counter, stepsize/telapsed));
  82. sys.stdout.flush();
  83. if self.lr_step == 'iteration': self.__scheduler__.step()
  84. if self.lr_step == 'epoch': self.__scheduler__.step()
  85. sys.stdout.write("\n");
  86. return (loss/counter, top1/counter);
  87. ## ===== ===== ===== ===== ===== ===== ===== =====
  88. ## Evaluate from list
  89. ## ===== ===== ===== ===== ===== ===== ===== =====
  90. def evaluateFromList(self, test_list, test_path, nDataLoaderThread, transform, print_interval=100, num_eval=10, **kwargs):
  91. self.__model__.eval();
  92. feats = {}
  93. tstart = time.time()
  94. ## Read all lines
  95. with open(test_list) as f:
  96. lines = f.readlines()
  97. ## Get a list of unique file names
  98. files = sum([x.strip().split(',')[-2:] for x in lines],[])
  99. setfiles = list(set(files))
  100. setfiles.sort()
  101. ## Define test data loader
  102. test_dataset = test_dataset_loader(setfiles, test_path, transform=transform, num_eval=num_eval, **kwargs)
  103. test_loader = torch.utils.data.DataLoader(
  104. test_dataset,
  105. batch_size=1,
  106. shuffle=False,
  107. num_workers=nDataLoaderThread,
  108. drop_last=False,
  109. )
  110. ## Extract features for every image
  111. for idx, data in enumerate(test_loader):
  112. inp1 = data[0][0].cuda()
  113. ref_feat = self.__model__(inp1).detach().cpu()
  114. feats[data[1][0]] = ref_feat
  115. telapsed = time.time() - tstart
  116. if idx % print_interval == 0:
  117. sys.stdout.write("\rReading %d of %d: %.2f Hz, embedding size %d"%(idx,len(setfiles),idx/telapsed,ref_feat.size()[1]));
  118. print('')
  119. all_scores = [];
  120. all_labels = [];
  121. tstart = time.time()
  122. ## Read files and compute all scores
  123. for idx, line in enumerate(lines):
  124. data = line.strip().split(',');
  125. ref_feat = feats[data[1]]
  126. com_feat = feats[data[2]]
  127. score = F.cosine_similarity(ref_feat, com_feat)
  128. all_scores.append(score);
  129. all_labels.append(int(data[0]));
  130. if idx % print_interval == 0:
  131. telapsed = time.time() - tstart
  132. sys.stdout.write("\rComputing %d of %d: %.2f Hz"%(idx,len(lines),idx/telapsed));
  133. sys.stdout.flush();
  134. print('')
  135. return (all_scores, all_labels);
  136. ## ===== ===== ===== ===== ===== ===== ===== =====
  137. ## Save parameters
  138. ## ===== ===== ===== ===== ===== ===== ===== =====
  139. def saveParameters(self, path):
  140. torch.save(self.__model__.state_dict(), path);
  141. ## ===== ===== ===== ===== ===== ===== ===== =====
  142. ## Load parameters
  143. ## ===== ===== ===== ===== ===== ===== ===== =====
  144. def loadParameters(self, path):
  145. self_state = self.__model__.state_dict();
  146. loaded_state = torch.load(path);
  147. for name, param in loaded_state.items():
  148. origname = name;
  149. if name not in self_state:
  150. if name not in self_state:
  151. print("%s is not in the model."%origname);
  152. continue;
  153. if self_state[name].size() != loaded_state[origname].size():
  154. print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, self_state[name].size(), loaded_state[origname].size()));
  155. continue;
  156. self_state[name].copy_(param);