import os import shutil import time import torch.backends.cudnn as cudnn import torch.nn.parallel import torch.optim import torch.utils.data # import torchnet as tnt import torchvision.transforms as transforms import torch.nn as nn import math import numpy as np # from util import * # tqdm.monitor_interval = 0 # class Engine(object): # def __init__(self, state={}): # self.state = state # if self._state('use_gpu') is None: # self.state['use_gpu'] = torch.cuda.is_available() # if self._state('image_size') is None: # self.state['image_size'] = 224 # if self._state('batch_size') is None: # self.state['batch_size'] = 64 # if self._state('workers') is None: # self.state['workers'] = 25 # if self._state('device_ids') is None: # self.state['device_ids'] = None # if self._state('evaluate') is None: # self.state['evaluate'] = False # if self._state('start_epoch') is None: # self.state['start_epoch'] = 0 # if self._state('max_epochs') is None: # self.state['max_epochs'] = 90 # if self._state('epoch_step') is None: # self.state['epoch_step'] = [] # # meters # self.state['meter_loss'] = tnt.meter.AverageValueMeter() # # time measure # self.state['batch_time'] = tnt.meter.AverageValueMeter() # self.state['data_time'] = tnt.meter.AverageValueMeter() # # display parameters # if self._state('use_pb') is None: # self.state['use_pb'] = False # if self._state('print_freq') is None: # self.state['print_freq'] = 10 # if not self.state['evaluate']: # self.state['log_path'] = open(self.state['save_model_path']+'train.log', 'a') # else: # self.state['log_path'] = open(self.state['save_model_path']+'eval.log', 'a') # def _state(self, name): # if name in self.state: # return self.state[name] # def init_learning(self, model, criterion): # if self._state('train_transform') is None: # normalize = transforms.Normalize(mean=model.image_normalization_mean, # std=model.image_normalization_std) # #flip = transforms.RandomChoice([transforms.RandomHorizontalFlip(),\ # #transforms.RandomVerticalFlip()]) # self.state['train_transform'] = transforms.Compose([ # transforms.Resize((512, 512)), # MultiScaleCrop(self.state['image_size'], scales=(1.0, 0.875, 0.75, 0.66, 0.5), max_distort=2), # transforms.RandomHorizontalFlip(), # transforms.ToTensor(), # normalize, # transforms.RandomErasing(p=0.7), # ]) # print(self.state['train_transform']) # if self._state('val_transform') is None: # normalize = transforms.Normalize(mean=model.image_normalization_mean, # std=model.image_normalization_std) # self.state['val_transform'] = transforms.Compose([ # Warp(self.state['image_size']), # transforms.ToTensor(), # normalize, # ]) # self.state['best_score'] = 0 # def learning(self, model, criterion, train_dataset, val_dataset, optimizer=None): # self.init_learning(model, criterion) # # define train and val transform # train_dataset.transform = self.state['train_transform'] # train_dataset.target_transform = self._state('train_target_transform') # val_dataset.transform = self.state['val_transform'] # val_dataset.target_transform = self._state('val_target_transform') # # data loading code # train_loader = torch.utils.data.DataLoader(train_dataset, # batch_size=self.state['batch_size'], shuffle=True, # num_workers=self.state['workers'],drop_last=True) # val_loader = torch.utils.data.DataLoader(val_dataset, # batch_size=self.state['batch_size'], shuffle=False, # num_workers=self.state['workers']) # # optionally resume from a checkpoint # if self._state('resume') is not None: # if os.path.isfile(self.state['resume']): # print("=> loading checkpoint '{}'".format(self.state['resume'])) # checkpoint = torch.load(self.state['resume']) # self.state['start_epoch'] = checkpoint['epoch'] # self.state['best_score'] = checkpoint['best_score'] # model.load_state_dict(checkpoint['state_dict']) # print("=> loaded checkpoint '{}' (epoch {})" # .format(self.state['evaluate'], checkpoint['epoch'])) # else: # print("=> no checkpoint found at '{}'".format(self.state['resume'])) # if self.state['use_gpu']: # train_loader.pin_memory = True # val_loader.pin_memory = True # cudnn.benchmark = True # model = torch.nn.DataParallel(model, device_ids=self.state['device_ids']).cuda() # criterion = criterion.cuda() # if self.state['evaluate']: # with torch.no_grad(): # self.validate(val_loader, model, criterion) # return # # TODO define optimizer # for epoch in range(self.state['start_epoch'], self.state['max_epochs']): # self.state['epoch'] = epoch # lr = self.adjust_learning_rate(optimizer) # print('lr:{:.5f}'.format(lr)) # # train for one epoch # self.train(train_loader, model, criterion, optimizer, epoch) # # evaluate on validation set # with torch.no_grad(): # prec1 = self.validate(val_loader, model, criterion) # # remember best prec@1 and save checkpoint # is_best = prec1 > self.state['best_score'] # self.state['best_score'] = max(prec1, self.state['best_score']) # self.save_checkpoint({ # 'epoch': epoch + 1, # 'arch': self._state('arch'), # 'state_dict': model.module.state_dict() if self.state['use_gpu'] else model.state_dict(), # 'best_score': self.state['best_score'], # }, is_best) # print(' *** best={best:.3f}'.format(best=self.state['best_score'])) # self.state['log_path'].close() # return self.state['best_score'] # def train(self, data_loader, model, criterion, optimizer, epoch): # # switch to train mode # model.train() # self.on_start_epoch(True, model, criterion, data_loader, optimizer) # if self.state['use_pb']: # data_loader = tqdm(data_loader, desc='Training') # end = time.time() # for i, (input, target) in enumerate(data_loader): # # measure data loading time # self.state['iteration'] = i # self.state['data_time_batch'] = time.time() - end # self.state['data_time'].add(self.state['data_time_batch']) # self.state['input'] = input # self.state['target'] = target # self.on_start_batch(True, model, criterion, data_loader, optimizer) # if self.state['use_gpu']: # self.state['target'] = self.state['target'].cuda() # self.on_forward(True, model, criterion, data_loader, optimizer) # # measure elapsed time # self.state['batch_time_current'] = time.time() - end # self.state['batch_time'].add(self.state['batch_time_current']) # end = time.time() # # measure accuracy # self.on_end_batch(True, model, criterion, data_loader, optimizer) # self.on_end_epoch(True, model, criterion, data_loader, optimizer) # def validate(self, data_loader, model, criterion): # # switch to evaluate mode # model.eval() # self.on_start_epoch(False, model, criterion, data_loader) # if self.state['use_pb']: # data_loader = tqdm(data_loader, desc='Test') # end = time.time() # for i, (input, target) in enumerate(data_loader): # # measure data loading time # self.state['iteration'] = i # self.state['data_time_batch'] = time.time() - end # self.state['data_time'].add(self.state['data_time_batch']) # self.state['input'] = input # self.state['target'] = target # self.on_start_batch(False, model, criterion, data_loader) # if self.state['use_gpu']: # self.state['target'] = self.state['target'].cuda() # self.on_forward(False, model, criterion, data_loader) # # measure elapsed time # self.state['batch_time_current'] = time.time() - end # self.state['batch_time'].add(self.state['batch_time_current']) # end = time.time() # # measure accuracy # self.on_end_batch(False, model, criterion, data_loader) # score = self.on_end_epoch(False, model, criterion, data_loader) # return score # def save_checkpoint(self, state, is_best, filename='checkpoint.pth.tar'): # if self._state('save_model_path') is not None: # filename_ = filename # filename = os.path.join(self.state['save_model_path'], filename_) # if not os.path.exists(self.state['save_model_path']): # os.makedirs(self.state['save_model_path']) # print('save model {filename}'.format(filename=filename)) # torch.save(state, filename) # if is_best: # filename_best = 'model_best.pth.tar' # if self._state('save_model_path') is not None: # filename_best = os.path.join(self.state['save_model_path'], filename_best) # shutil.copyfile(filename, filename_best) # if self._state('save_model_path') is not None: # if self._state('filename_previous_best') is not None: # os.remove(self._state('filename_previous_best')) # filename_best = os.path.join(self.state['save_model_path'], 'model_best_{score:.4f}.pth.tar'.format(score=state['best_score'])) # shutil.copyfile(filename, filename_best) # self.state['filename_previous_best'] = filename_best # def adjust_learning_rate(self, optimizer): # """Sets the learning rate to the initial LR decayed by 10 every 30 epochs""" # # lr = args.lr * (0.1 ** (epoch // 30)) # decay = 0.1 ** (sum(self.state['epoch'] >= np.array(self.state['epoch_step']))) # lr = self.state['lr'] * decay # for param_group in optimizer.param_groups: # param_group['lr'] = lr # return lr # class MultiLabelEngine(Engine): # def __init__(self, state): # Engine.__init__(self, state) # if self._state('difficult_examples') is None: # self.state['difficult_examples'] = False # self.state['ap_meter'] = AveragePrecisionMeter(self.state['difficult_examples']) # def on_start_epoch(self, training, model, criterion, data_loader, optimizer=None, display=True): # self.state['meter_loss'].reset() # self.state['batch_time'].reset() # self.state['data_time'].reset() # self.state['ap_meter'].reset() # def on_end_epoch(self, training, model, criterion, data_loader, optimizer=None, display=True): # map = 100 * self.state['ap_meter'].value().mean() # loss = self.state['meter_loss'].value()[0] # OP, OR, OF1, CP, CR, CF1 = self.state['ap_meter'].overall() # OP_k, OR_k, OF1_k, CP_k, CR_k, CF1_k = self.state['ap_meter'].overall_topk(3) # if display: # if training: # print('Epoch: [{0}]\t' # 'Loss {loss:.4f}\t' # 'mAP {map:.3f}'.format(self.state['epoch'], loss=loss, map=map)) # print('OP: {OP:.4f}\t' # 'OR: {OR:.4f}\t' # 'OF1: {OF1:.4f}\t' # 'CP: {CP:.4f}\t' # 'CR: {CR:.4f}\t' # 'CF1: {CF1:.4f}'.format(OP=OP, OR=OR, OF1=OF1, CP=CP, CR=CR, CF1=CF1)) # self.state['log_path'].write('Epoch: [{0}]\t' # 'Loss {loss:.4f}\t' # 'mAP {map:.3f}'.format(self.state['epoch'], loss=loss, map=map)+'\n') # self.state['log_path'].write('OP: {OP:.4f}\t' # 'OR: {OR:.4f}\t' # 'OF1: {OF1:.4f}\t' # 'CP: {CP:.4f}\t' # 'CR: {CR:.4f}\t' # 'CF1: {CF1:.4f}'.format(OP=OP, OR=OR, OF1=OF1, CP=CP, CR=CR, CF1=CF1)+'\n') # else: # print('Test: \t Loss {loss:.4f}\t mAP {map:.3f}'.format(loss=loss, map=map)) # print('OP: {OP:.4f}\t' # 'OR: {OR:.4f}\t' # 'OF1: {OF1:.4f}\t' # 'CP: {CP:.4f}\t' # 'CR: {CR:.4f}\t' # 'CF1: {CF1:.4f}'.format(OP=OP, OR=OR, OF1=OF1, CP=CP, CR=CR, CF1=CF1)) # print('OP_3: {OP:.4f}\t' # 'OR_3: {OR:.4f}\t' # 'OF1_3: {OF1:.4f}\t' # 'CP_3: {CP:.4f}\t' # 'CR_3: {CR:.4f}\t' # 'CF1_3: {CF1:.4f}'.format(OP=OP_k, OR=OR_k, OF1=OF1_k, CP=CP_k, CR=CR_k, CF1=CF1_k)) # self.state['log_path'].write('Test: \t Loss {loss:.4f}\t mAP {map:.3f}'.format(loss=loss, map=map)+'\n') # self.state['log_path'].write('OP: {OP:.4f}\t' # 'OR: {OR:.4f}\t' # 'OF1: {OF1:.4f}\t' # 'CP: {CP:.4f}\t' # 'CR: {CR:.4f}\t' # 'CF1: {CF1:.4f}'.format(OP=OP, OR=OR, OF1=OF1, CP=CP, CR=CR, CF1=CF1)+'\n') # self.state['log_path'].write('OP_3: {OP:.4f}\t' # 'OR_3: {OR:.4f}\t' # 'OF1_3: {OF1:.4f}\t' # 'CP_3: {CP:.4f}\t' # 'CR_3: {CR:.4f}\t' # 'CF1_3: {CF1:.4f}'.format(OP=OP_k, OR=OR_k, OF1=OF1_k, CP=CP_k, CR=CR_k, CF1=CF1_k)+'\n') # return map # def on_start_batch(self, training, model, criterion, data_loader, optimizer=None, display=True): # self.state['target_gt'] = self.state['target'].clone() # self.state['target'][self.state['target'] == 0] = 1 # self.state['target'][self.state['target'] == -1] = 0 # input = self.state['input'] # self.state['feature'] = input[0] # self.state['out'] = input[1] # def on_end_batch(self, training, model, criterion, data_loader, optimizer=None, display=True): # self.state['loss_batch'] = self.state['loss'].item() # self.state['meter_loss'].add(self.state['loss_batch']) # # measure mAP # self.state['ap_meter'].add(self.state['output'].data, self.state['target_gt']) # if display and self.state['print_freq'] != 0 and self.state['iteration'] % self.state['print_freq'] == 0: # loss = self.state['meter_loss'].value()[0] # batch_time = self.state['batch_time'].value()[0] # data_time = self.state['data_time'].value()[0] # if training: # print('Epoch: [{0}][{1}/{2}]\t' # 'Time {batch_time_current:.3f} ({batch_time:.3f})\t' # 'Data {data_time_current:.3f} ({data_time:.3f})\t' # 'Loss {loss_current:.4f} ({loss:.4f})'.format( # self.state['epoch'], self.state['iteration'], len(data_loader), # batch_time_current=self.state['batch_time_current'], # batch_time=batch_time, data_time_current=self.state['data_time_batch'], # data_time=data_time, loss_current=self.state['loss_batch'], loss=loss)) # else: # print('Test: [{0}/{1}]\t' # 'Time {batch_time_current:.3f} ({batch_time:.3f})\t' # 'Data {data_time_current:.3f} ({data_time:.3f})\t' # 'Loss {loss_current:.4f} ({loss:.4f})'.format( # self.state['iteration'], len(data_loader), batch_time_current=self.state['batch_time_current'], # batch_time=batch_time, data_time_current=self.state['data_time_batch'], # data_time=data_time, loss_current=self.state['loss_batch'], loss=loss)) # def on_forward(self, training, model, criterion, data_loader, optimizer=None, display=True): # feature_var = self.state['feature'] # target_var = self.state['target'] # #print(type(feature_var),feature_var.size()) # # compute output # self.state['output'] = model(feature_var) # self.state['loss'] = criterion(self.state['output'], target_var) # if training: # optimizer.zero_grad() # self.state['loss'].backward() # nn.utils.clip_grad_norm(model.parameters(), max_norm=10.0) # optimizer.step() class AveragePrecisionMeter(object): """ The APMeter measures the average precision per class. The APMeter is designed to operate on `NxK` Tensors `output` and `target`, and optionally a `Nx1` Tensor weight where (1) the `output` contains model output scores for `N` examples and `K` classes that ought to be higher when the model is more convinced that the example should be positively labeled, and smaller when the model believes the example should be negatively labeled (for instance, the output of a sigmoid function); (2) the `target` contains only values 0 (for negative examples) and 1 (for positive examples); and (3) the `weight` ( > 0) represents weight for each sample. """ def __init__(self, difficult_examples=False): super(AveragePrecisionMeter, self).__init__() self.reset() self.difficult_examples = difficult_examples def reset(self): """Resets the meter with empty member variables""" self.scores = torch.FloatTensor(torch.FloatStorage()) self.targets = torch.LongTensor(torch.LongStorage()) def add(self, output, target): """ Args: output (Tensor): NxK tensor that for each of the N examples indicates the probability of the example belonging to each of the K classes, according to the model. The probabilities should sum to one over all classes target (Tensor): binary NxK tensort that encodes which of the K classes are associated with the N-th input (eg: a row [0, 1, 0, 1] indicates that the example is associated with classes 2 and 4) weight (optional, Tensor): Nx1 tensor representing the weight for each example (each weight > 0) """ if not torch.is_tensor(output): output = torch.from_numpy(output) if not torch.is_tensor(target): target = torch.from_numpy(target) if output.dim() == 1: output = output.view(-1, 1) else: assert output.dim() == 2, \ 'wrong output size (should be 1D or 2D with one column \ per class)' if target.dim() == 1: target = target.view(-1, 1) else: assert target.dim() == 2, \ 'wrong target size (should be 1D or 2D with one column \ per class)' if self.scores.numel() > 0: assert target.size(1) == self.targets.size(1), \ 'dimensions for output should match previously added examples.' # make sure storage is of sufficient size if self.scores.storage().size() < self.scores.numel() + output.numel(): new_size = math.ceil(self.scores.storage().size() * 1.5) self.scores.storage().resize_(int(new_size + output.numel())) self.targets.storage().resize_(int(new_size + output.numel())) # store scores and targets offset = self.scores.size(0) if self.scores.dim() > 0 else 0 self.scores.resize_(offset + output.size(0), output.size(1)) self.targets.resize_(offset + target.size(0), target.size(1)) self.scores.narrow(0, offset, output.size(0)).copy_(output) self.targets.narrow(0, offset, target.size(0)).copy_(target) def value(self): """Returns the model's average precision for each class Return: ap (FloatTensor): 1xK tensor, with avg precision for each class k """ if self.scores.numel() == 0: return 0 ap = torch.zeros(self.scores.size(1)) rg = torch.arange(1, self.scores.size(0)).float() # self.save_to_mat() # self.save_to_np() # compute average precision for each class for k in range(self.scores.size(1)): # sort scores scores = self.scores[:, k] targets = self.targets[:, k] # compute average precision ap[k] = AveragePrecisionMeter.average_precision(scores, targets, self.difficult_examples) return ap @staticmethod def average_precision(output, target, difficult_examples=True): if target.sum() == 0: return 0 else: # sort examples sorted, indices = torch.sort(output, dim=0, descending=True) # Computes prec@i pos_count = 0. total_count = 0. precision_at_i = 0. for i in indices: label = target[i] if difficult_examples and label == 0: continue if label == 1: pos_count += 1 total_count += 1 if label == 1: precision_at_i += pos_count / total_count precision_at_i /= pos_count return precision_at_i def overall(self): if self.scores.numel() == 0: return 0 scores = self.scores.cpu().numpy() targets = self.targets.cpu().numpy() targets[targets == -1] = 0 return self.evaluation(scores, targets) def overall_topk(self, k): targets = self.targets.cpu().numpy() targets[targets == -1] = 0 n, c = self.scores.size() scores = np.zeros((n, c)) - 1 index = self.scores.topk(k, 1, True, True)[1].cpu().numpy() tmp = self.scores.cpu().numpy() for i in range(n): for ind in index[i]: scores[i, ind] = 1 if tmp[i, ind] >= 0 else -1 return self.evaluation(scores, targets) def evaluation(self, scores_, targets_): n, n_class = scores_.shape Nc, Np, Ng = np.zeros(n_class), np.zeros(n_class), np.zeros(n_class) for k in range(n_class): scores = scores_[:, k] targets = targets_[:, k] targets[targets == -1] = 0 Ng[k] = np.sum(targets == 1) Np[k] = np.sum(scores >= 0) Nc[k] = np.sum(targets * (scores >= 0)) Np[Np == 0] = 1 OP = np.sum(Nc) / np.sum(Np) OR = np.sum(Nc) / np.sum(Ng) OF1 = (2 * OP * OR) / (OP + OR) CP = np.sum(Nc / Np) / n_class CR = np.sum(Nc / Ng) / n_class CF1 = (2 * CP * CR) / (CP + CR) return OP, OR, OF1, CP, CR, CF1