captchaboy/dfff4444
0
1import logging2import shutil3import time4 5import editdistance as ed6import torchvision.utils as vutils7from fastai.callbacks.tensorboard import (LearnerTensorboardWriter,8 SummaryWriter, TBWriteRequest,9 asyncTBWriter)10from fastai.vision import *11from torch.nn.parallel import DistributedDataParallel12from torchvision import transforms13 14import dataset15from utils import CharsetMapper, Timer, blend_mask16 17 18class IterationCallback(LearnerTensorboardWriter):19 "A `TrackerCallback` that monitor in each iteration."20 def __init__(self, learn:Learner, name:str='model', checpoint_keep_num=5,21 show_iters:int=50, eval_iters:int=1000, save_iters:int=20000,22 start_iters:int=0, stats_iters=20000):23 #if self.learn.rank is not None: time.sleep(self.learn.rank) # keep all event files24 super().__init__(learn, base_dir='.', name=learn.path, loss_iters=show_iters, 25 stats_iters=stats_iters, hist_iters=stats_iters)26 self.name, self.bestname = Path(name).name, f'best-{Path(name).name}'27 self.show_iters = show_iters28 self.eval_iters = eval_iters29 self.save_iters = save_iters30 self.start_iters = start_iters31 self.checpoint_keep_num = checpoint_keep_num32 self.metrics_root = 'metrics/' # rewrite33 self.timer = Timer()34 self.host = self.learn.rank is None or self.learn.rank == 035 36 def _write_metrics(self, iteration:int, names:List[str], last_metrics:MetricsList)->None:37 "Writes training metrics to Tensorboard."38 for i, name in enumerate(names):39 if last_metrics is None or len(last_metrics) < i+1: return40 scalar_value = last_metrics[i]41 self._write_scalar(name=name, scalar_value=scalar_value, iteration=iteration)42 43 def _write_sub_loss(self, iteration:int, last_losses:dict)->None:44 "Writes sub loss to Tensorboard."45 for name, loss in last_losses.items():46 scalar_value = to_np(loss)47 tag = self.metrics_root + name48 self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)49 50 def _save(self, name):51 if isinstance(self.learn.model, DistributedDataParallel):52 tmp = self.learn.model53 self.learn.model = self.learn.model.module54 self.learn.save(name)55 self.learn.model = tmp56 else: self.learn.save(name)57 58 def _validate(self, dl=None, callbacks=None, metrics=None, keeped_items=False):59 "Validate on `dl` with potential `callbacks` and `metrics`."60 dl = ifnone(dl, self.learn.data.valid_dl)61 metrics = ifnone(metrics, self.learn.metrics)62 cb_handler = CallbackHandler(ifnone(callbacks, []), metrics)63 cb_handler.on_train_begin(1, None, metrics); cb_handler.on_epoch_begin()64 if keeped_items: cb_handler.state_dict.update(dict(keeped_items=[]))65 val_metrics = validate(self.learn.model, dl, self.loss_func, cb_handler)66 cb_handler.on_epoch_end(val_metrics)67 if keeped_items: return cb_handler.state_dict['keeped_items']68 else: return cb_handler.state_dict['last_metrics']69 70 def jump_to_epoch_iter(self, epoch:int, iteration:int)->None:71 try:72 self.learn.load(f'{self.name}_{epoch}_{iteration}', purge=False)73 logging.info(f'Loaded {self.name}_{epoch}_{iteration}')74 except: logging.info(f'Model {self.name}_{epoch}_{iteration} not found.')75 76 def on_train_begin(self, n_epochs, **kwargs):77 # TODO: can not write graph here78 # super().on_train_begin(**kwargs)79 self.best = -float('inf')80 self.timer.tic()81 if self.host:82 checkpoint_path = self.learn.path/'checkpoint.yaml'83 if checkpoint_path.exists():84 os.remove(checkpoint_path)85 open(checkpoint_path, 'w').close()86 return {'skip_validate': True, 'iteration':self.start_iters} # disable default validate87 88 def on_batch_begin(self, **kwargs:Any)->None:89 self.timer.toc_data()90 super().on_batch_begin(**kwargs)91 92 def on_batch_end(self, iteration, epoch, last_loss, smooth_loss, train, **kwargs):93 super().on_batch_end(last_loss, iteration, train, **kwargs)94 if iteration == 0: return95 96 if iteration % self.loss_iters == 0:97 last_losses = self.learn.loss_func.last_losses98 self._write_sub_loss(iteration=iteration, last_losses=last_losses)99 self.tbwriter.add_scalar(tag=self.metrics_root + 'lr',100 scalar_value=self.opt.lr, global_step=iteration)101 102 if iteration % self.show_iters == 0:103 log_str = f'epoch {epoch} iter {iteration}: loss = {last_loss:6.4f}, ' \104 f'smooth loss = {smooth_loss:6.4f}'105 logging.info(log_str)106 # log_str = f'data time = {self.timer.data_diff:.4f}s, runing time = {self.timer.running_diff:.4f}s'107 # logging.info(log_str)108 109 if iteration % self.eval_iters == 0:110 # TODO: or remove time to on_epoch_end111 # 1. Record time 112 log_str = f'average data time = {self.timer.average_data_time():.4f}s, ' \113 f'average running time = {self.timer.average_running_time():.4f}s'114 logging.info(log_str)115 116 # 2. Call validate117 last_metrics = self._validate()118 self.learn.model.train()119 log_str = f'epoch {epoch} iter {iteration}: eval loss = {last_metrics[0]:6.4f}, ' \120 f'ccr = {last_metrics[1]:6.4f}, cwr = {last_metrics[2]:6.4f}, ' \121 f'ted = {last_metrics[3]:6.4f}, ned = {last_metrics[4]:6.4f}, ' \122 f'ted/w = {last_metrics[5]:6.4f}, '123 logging.info(log_str)124 names = ['eval_loss', 'ccr', 'cwr', 'ted', 'ned', 'ted/w']125 self._write_metrics(iteration, names, last_metrics)126 127 # 3. Save best model128 current = last_metrics[2]129 if current is not None and current > self.best:130 logging.info(f'Better model found at epoch {epoch}, '\131 f'iter {iteration} with accuracy value: {current:6.4f}.')132 self.best = current133 self._save(f'{self.bestname}')134 135 if iteration % self.save_iters == 0 and self.host:136 logging.info(f'Save model {self.name}_{epoch}_{iteration}')137 filename = f'{self.name}_{epoch}_{iteration}'138 self._save(filename)139 140 checkpoint_path = self.learn.path/'checkpoint.yaml'141 if not checkpoint_path.exists():142 open(checkpoint_path, 'w').close()143 with open(checkpoint_path, 'r') as file:144 checkpoints = yaml.load(file, Loader=yaml.FullLoader) or dict()145 checkpoints['all_checkpoints'] = (146 checkpoints.get('all_checkpoints') or list())147 checkpoints['all_checkpoints'].insert(0, filename)148 if len(checkpoints['all_checkpoints']) > self.checpoint_keep_num:149 removed_checkpoint = checkpoints['all_checkpoints'].pop()150 removed_checkpoint = self.learn.path/self.learn.model_dir/f'{removed_checkpoint}.pth'151 os.remove(removed_checkpoint)152 checkpoints['current_checkpoint'] = filename153 with open(checkpoint_path, 'w') as file:154 yaml.dump(checkpoints, file)155 156 157 self.timer.toc_running()158 159 def on_train_end(self, **kwargs):160 #self.learn.load(f'{self.bestname}', purge=False)161 pass162 163 def on_epoch_end(self, last_metrics:MetricsList, iteration:int, **kwargs)->None:164 self._write_embedding(iteration=iteration)165 166 167class TextAccuracy(Callback):168 _names = ['ccr', 'cwr', 'ted', 'ned', 'ted/w']169 def __init__(self, charset_path, max_length, case_sensitive, model_eval):170 self.charset_path = charset_path171 self.max_length = max_length172 self.case_sensitive = case_sensitive173 self.charset = CharsetMapper(charset_path, self.max_length)174 self.names = self._names175 176 self.model_eval = model_eval or 'alignment'177 assert self.model_eval in ['vision', 'language', 'alignment']178 179 def on_epoch_begin(self, **kwargs):180 self.total_num_char = 0.181 self.total_num_word = 0.182 self.correct_num_char = 0.183 self.correct_num_word = 0.184 self.total_ed = 0.185 self.total_ned = 0.186 187 def _get_output(self, last_output):188 if isinstance(last_output, (tuple, list)): 189 for res in last_output:190 if res['name'] == self.model_eval: output = res191 else: output = last_output192 return output193 194 def _update_output(self, last_output, items):195 if isinstance(last_output, (tuple, list)): 196 for res in last_output:197 if res['name'] == self.model_eval: res.update(items)198 else: last_output.update(items)199 return last_output200 201 def on_batch_end(self, last_output, last_target, **kwargs):202 output = self._get_output(last_output)203 logits, pt_lengths = output['logits'], output['pt_lengths']204 pt_text, pt_scores, pt_lengths_ = self.decode(logits)205 assert (pt_lengths == pt_lengths_).all(), f'{pt_lengths} != {pt_lengths_} for {pt_text}'206 last_output = self._update_output(last_output, {'pt_text':pt_text, 'pt_scores':pt_scores})207 208 pt_text = [self.charset.trim(t) for t in pt_text]209 label = last_target[0]210 if label.dim() == 3: label = label.argmax(dim=-1) # one-hot label211 gt_text = [self.charset.get_text(l, trim=True) for l in label]212 213 for i in range(len(gt_text)):214 if not self.case_sensitive:215 gt_text[i], pt_text[i] = gt_text[i].lower(), pt_text[i].lower()216 distance = ed.eval(gt_text[i], pt_text[i])217 self.total_ed += distance218 self.total_ned += float(distance) / max(len(gt_text[i]), 1)219 220 if gt_text[i] == pt_text[i]:221 self.correct_num_word += 1222 self.total_num_word += 1223 224 for j in range(min(len(gt_text[i]), len(pt_text[i]))):225 if gt_text[i][j] == pt_text[i][j]:226 self.correct_num_char += 1227 self.total_num_char += len(gt_text[i])228 229 return {'last_output': last_output}230 231 def on_epoch_end(self, last_metrics, **kwargs):232 mets = [self.correct_num_char / self.total_num_char,233 self.correct_num_word / self.total_num_word,234 self.total_ed,235 self.total_ned,236 self.total_ed / self.total_num_word]237 return add_metrics(last_metrics, mets)238 239 def decode(self, logit):240 """ Greed decode """241 # TODO: test running time and decode on GPU242 out = F.softmax(logit, dim=2)243 pt_text, pt_scores, pt_lengths = [], [], []244 for o in out:245 text = self.charset.get_text(o.argmax(dim=1), padding=False, trim=False)246 text = text.split(self.charset.null_char)[0] # end at end-token247 pt_text.append(text)248 pt_scores.append(o.max(dim=1)[0])249 pt_lengths.append(min(len(text) + 1, self.max_length)) # one for end-token250 pt_scores = torch.stack(pt_scores)251 pt_lengths = pt_scores.new_tensor(pt_lengths, dtype=torch.long)252 return pt_text, pt_scores, pt_lengths253 254 255class TopKTextAccuracy(TextAccuracy):256 _names = ['ccr', 'cwr']257 def __init__(self, k, charset_path, max_length, case_sensitive, model_eval):258 self.k = k259 self.charset_path = charset_path260 self.max_length = max_length261 self.case_sensitive = case_sensitive262 self.charset = CharsetMapper(charset_path, self.max_length)263 self.names = self._names264 265 def on_epoch_begin(self, **kwargs):266 self.total_num_char = 0.267 self.total_num_word = 0.268 self.correct_num_char = 0.269 self.correct_num_word = 0.270 271 def on_batch_end(self, last_output, last_target, **kwargs):272 logits, pt_lengths = last_output['logits'], last_output['pt_lengths']273 gt_labels, gt_lengths = last_target[:]274 275 for logit, pt_length, label, length in zip(logits, pt_lengths, gt_labels, gt_lengths):276 word_flag = True277 for i in range(length):278 char_logit = logit[i].topk(self.k)[1]279 char_label = label[i].argmax(-1)280 if char_label in char_logit: self.correct_num_char += 1281 else: word_flag = False282 self.total_num_char += 1283 if pt_length == length and word_flag:284 self.correct_num_word += 1285 self.total_num_word += 1286 287 def on_epoch_end(self, last_metrics, **kwargs):288 mets = [self.correct_num_char / self.total_num_char,289 self.correct_num_word / self.total_num_word,290 0., 0., 0.]291 return add_metrics(last_metrics, mets)292 293 294class DumpPrediction(LearnerCallback):295 296 def __init__(self, learn, dataset, charset_path, model_eval, image_only=False, debug=False):297 super().__init__(learn=learn)298 self.debug = debug299 self.model_eval = model_eval or 'alignment'300 self.image_only = image_only301 assert self.model_eval in ['vision', 'language', 'alignment']302 303 self.dataset, self.root = dataset, Path(self.learn.path)/f'{dataset}-{self.model_eval}'304 self.attn_root = self.root/'attn'305 self.charset = CharsetMapper(charset_path)306 if self.root.exists(): shutil.rmtree(self.root)307 self.root.mkdir(), self.attn_root.mkdir()308 309 self.pil = transforms.ToPILImage()310 self.tensor = transforms.ToTensor()311 size = self.learn.data.img_h, self.learn.data.img_w312 self.resize = transforms.Resize(size=size, interpolation=0)313 self.c = 0314 315 def on_batch_end(self, last_input, last_output, last_target, **kwargs):316 if isinstance(last_output, (tuple, list)):317 for res in last_output:318 if res['name'] == self.model_eval: pt_text = res['pt_text']319 if res['name'] == 'vision': attn_scores = res['attn_scores'].detach().cpu()320 if res['name'] == self.model_eval: logits = res['logits']321 else:322 pt_text = last_output['pt_text']323 attn_scores = last_output['attn_scores'].detach().cpu()324 logits = last_output['logits']325 326 images = last_input[0] if isinstance(last_input, (tuple, list)) else last_input327 images = images.detach().cpu()328 pt_text = [self.charset.trim(t) for t in pt_text]329 gt_label = last_target[0]330 if gt_label.dim() == 3: gt_label = gt_label.argmax(dim=-1) # one-hot label331 gt_text = [self.charset.get_text(l, trim=True) for l in gt_label]332 333 prediction, false_prediction = [], []334 for gt, pt, image, attn, logit in zip(gt_text, pt_text, images, attn_scores, logits):335 prediction.append(f'{gt}\t{pt}\n')336 if gt != pt:337 if self.debug:338 scores = torch.softmax(logit, dim=-1)[:max(len(pt), len(gt)) + 1]339 logging.info(f'{self.c} gt {gt}, pt {pt}, logit {logit.shape}, scores {scores.topk(5, dim=-1)}')340 false_prediction.append(f'{gt}\t{pt}\n')341 342 image = self.learn.data.denorm(image)343 if not self.image_only:344 image_np = np.array(self.pil(image))345 attn_pil = [self.pil(a) for a in attn[:, None, :, :]]346 attn = [self.tensor(self.resize(a)).repeat(3, 1, 1) for a in attn_pil]347 attn_sum = np.array([np.array(a) for a in attn_pil[:len(pt)]]).sum(axis=0)348 blended_sum = self.tensor(blend_mask(image_np, attn_sum))349 blended = [self.tensor(blend_mask(image_np, np.array(a))) for a in attn_pil]350 save_image = torch.stack([image] + attn + [blended_sum] + blended)351 save_image = save_image.view(2, -1, *save_image.shape[1:])352 save_image = save_image.permute(1, 0, 2, 3, 4).flatten(0, 1)353 vutils.save_image(save_image, self.attn_root/f'{self.c}_{gt}_{pt}.jpg', 354 nrow=2, normalize=True, scale_each=True)355 else:356 self.pil(image).save(self.attn_root/f'{self.c}_{gt}_{pt}.jpg')357 self.c += 1358 359 with open(self.root/f'{self.model_eval}.txt', 'a') as f: f.writelines(prediction)360 with open(self.root/f'{self.model_eval}-false.txt', 'a') as f: f.writelines(false_prediction) 361