123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222 |
- #!/usr/bin/python
- #-*- coding: utf-8 -*-
- import torch, pickle
- import torch.nn as nn
- import torch.nn.functional as F
- import numpy, math, pdb, sys
- import time, importlib
- from DatasetLoader import test_dataset_loader
- from torch.cuda.amp import autocast, GradScaler
- import cv2
- import glob
- from PIL import Image
- from flask import Flask, request
- class EmbedNet(nn.Module):
- def __init__(self, model, optimizer, trainfunc, nPerClass, **kwargs):
- super(EmbedNet, self).__init__();
- ## __S__ is the embedding model
- EmbedNetModel = importlib.import_module('models.'+model).__getattribute__('MainModel')
- self.__S__ = EmbedNetModel(**kwargs);
- ## __L__ is the classifier plus the loss function
- LossFunction = importlib.import_module('loss.'+trainfunc).__getattribute__('LossFunction')
- self.__L__ = LossFunction(**kwargs);
- ## Number of examples per identity per batch
- self.nPerClass = nPerClass
- def forward(self, data, label=None):
- data = data.reshape(-1,data.size()[-3],data.size()[-2],data.size()[-1])
- outp = self.__S__.forward(data)
- if label == None:
- return outp
- else:
- outp = outp.reshape(self.nPerClass,-1,outp.size()[-1]).transpose(1,0).squeeze(1)
- nloss, prec1 = self.__L__.forward(outp,label)
- return nloss, prec1
- class ModelTrainer(object):
- def __init__(self, embed_model, optimizer, scheduler, mixedprec, **kwargs):
- self.__model__ = embed_model
- ## Optimizer (e.g. Adam or SGD)
- Optimizer = importlib.import_module('optimizer.'+optimizer).__getattribute__('Optimizer')
- self.__optimizer__ = Optimizer(self.__model__.parameters(), **kwargs)
- ## Learning rate scheduler
- Scheduler = importlib.import_module('scheduler.'+scheduler).__getattribute__('Scheduler')
- self.__scheduler__, self.lr_step = Scheduler(self.__optimizer__, **kwargs)
- ## For mixed precision training
- self.scaler = GradScaler()
- self.mixedprec = mixedprec
- assert self.lr_step in ['epoch', 'iteration']
- # ## ===== ===== ===== ===== ===== ===== ===== =====
- # ## Train network
- # ## ===== ===== ===== ===== ===== ===== ===== =====
- def train_network(self, loader, verbose):
- self.__model__.train();
- stepsize = loader.batch_size;
- counter = 0;
- index = 0;
- loss = 0;
- top1 = 0 # EER or accuracy
- tstart = time.time()
-
- for data, label in loader:
- data = data.transpose(1,0)
- ## Reset gradients
- self.__model__.zero_grad();
- ## Forward and backward passes
- if self.mixedprec:
- with autocast():
- nloss, prec1 = self.__model__(data.cuda(), label.cuda())
- self.scaler.scale(nloss).backward();
- self.scaler.step(self.__optimizer__);
- self.scaler.update();
- else:
- nloss, prec1 = self.__model__(data.cuda(), label.cuda())
- nloss.backward();
- self.__optimizer__.step();
- loss += nloss.detach().cpu();
- top1 += prec1.detach().cpu();
- counter += 1;
- index += stepsize;
- telapsed = time.time() - tstart
- tstart = time.time()
- if verbose:
- sys.stdout.write("\rProcessing (%d) "%(index));
- sys.stdout.write("Loss %f TEER/TAcc %2.3f%% - %.2f Hz "%(loss/counter, top1/counter, stepsize/telapsed));
- sys.stdout.flush();
- if self.lr_step == 'iteration': self.__scheduler__.step()
- if self.lr_step == 'epoch': self.__scheduler__.step()
- sys.stdout.write("\n");
-
- return (loss/counter, top1/counter);
- ## ===== ===== ===== ===== ===== ===== ===== =====
- ## Evaluate from list
- ## ===== ===== ===== ===== ===== ===== ===== =====
- def evaluateFromList(self, test_list, test_path, nDataLoaderThread, transform, print_interval=100, num_eval=10, **kwargs):
-
- self.__model__.eval();
-
- feats = {}
- tstart = time.time()
- ## Read all lines
- with open(test_list) as f:
- lines = f.readlines()
- ## Get a list of unique file names
- files = sum([x.strip().split(',')[-2:] for x in lines],[])
- setfiles = list(set(files))
- setfiles.sort()
- ## Define test data loader
- test_dataset = test_dataset_loader(setfiles, test_path, transform=transform, num_eval=num_eval, **kwargs)
- test_loader = torch.utils.data.DataLoader(
- test_dataset,
- batch_size=1,
- shuffle=False,
- num_workers=nDataLoaderThread,
- drop_last=False,
- )
- ## Extract features for every image
- for idx, data in enumerate(test_loader):
- inp1 = data[0][0].cuda()
- ref_feat = self.__model__(inp1).detach().cpu()
- feats[data[1][0]] = ref_feat
- telapsed = time.time() - tstart
- if idx % print_interval == 0:
- sys.stdout.write("\rReading %d of %d: %.2f Hz, embedding size %d"%(idx,len(setfiles),idx/telapsed,ref_feat.size()[1]));
- print('')
- all_scores = [];
- all_labels = [];
- tstart = time.time()
- ## Read files and compute all scores
- for idx, line in enumerate(lines):
- data = line.strip().split(',');
- ref_feat = feats[data[1]]
- com_feat = feats[data[2]]
- score = F.cosine_similarity(ref_feat, com_feat)
- all_scores.append(score);
- all_labels.append(int(data[0]));
- if idx % print_interval == 0:
- telapsed = time.time() - tstart
- sys.stdout.write("\rComputing %d of %d: %.2f Hz"%(idx,len(lines),idx/telapsed));
- sys.stdout.flush();
- print('')
- return (all_scores, all_labels);
- ## ===== ===== ===== ===== ===== ===== ===== =====
- ## Save parameters
- ## ===== ===== ===== ===== ===== ===== ===== =====
- def saveParameters(self, path):
-
- torch.save(self.__model__.state_dict(), path);
- ## ===== ===== ===== ===== ===== ===== ===== =====
- ## Load parameters
- ## ===== ===== ===== ===== ===== ===== ===== =====
- def loadParameters(self, path):
- self_state = self.__model__.state_dict();
- loaded_state = torch.load(path);
- for name, param in loaded_state.items():
- origname = name;
- if name not in self_state:
- if name not in self_state:
- print("%s is not in the model."%origname);
- continue;
- if self_state[name].size() != loaded_state[origname].size():
- print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, self_state[name].size(), loaded_state[origname].size()));
- continue;
- self_state[name].copy_(param);
|