def infer(loader, net, device): net.eval() metric = ConfusionMatrix(len(loader.dataset.classes) - 1) for i, (x, gt) in enumerate(loader): x = x.to(device, non_blocking=True) gt = gt.squeeze(1).to(device, dtype=torch.long, non_blocking=True) pred = net(x)["fine"].argmax(1) metric.update(pred, gt) mIoU = metric.mIoU() logging.info(f"[Infer] mIOU : {mIoU}") return mIoU
class Evaluator(object): # def __init__(self, n_class, size0, size_g, size_p, n, sub_batch_size=6, mode=1, test=False): def __init__(self, n_class, size_g, size_p, sub_batch_size=6, mode=1, test=False): self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class # self.size0 = size0 self.size_g = size_g self.size_p = size_p # self.n = n self.sub_batch_size = sub_batch_size self.mode = mode self.test = test # self.ratio = float(size_p[0]) / size0[0] # self.step = (size0[0] - size_p[0]) // (n - 1) # self.template, self.coordinates = template_patch2global(size0, size_p, n, self.step) if test: self.flip_range = [False, True] self.rotate_range = [0, 1, 2, 3] else: self.flip_range = [False] self.rotate_range = [0] def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def eval_test(self, sample, model, global_fixed): with torch.no_grad(): images = sample["image"] if not self.test: labels = sample["label"] # PIL images #lbls = [RGB_mapping_to_class(np.array(label)) for label in labels] #labels = [Image.fromarray(lbl) for lbl in lbls] labels_npy = masks_transform(labels, numpy=True) images_global = resize(images, self.size_g) outputs_global = np.zeros( (len(images), self.n_class, self.size_g[0] // 4, self.size_g[1] // 4) ) # outputs_global = np.zeros((len(images), self.n_class, self.size_g[0], self.size_g[1])) if self.mode == 2 or self.mode == 3: images_local = [image.copy() for image in images] # scores_local = np.zeros((len(images), self.n_class, self.size0[0], self.size0[1])) # scores = np.zeros((len(images), self.n_class, self.size0[0], self.size0[1])) scores_local = [ np.zeros((1, self.n_class, images[i].size[1], images[i].size[0])) for i in range(len(images)) ] scores = [ np.zeros((1, self.n_class, images[i].size[1], images[i].size[0])) for i in range(len(images)) ] for flip in self.flip_range: if flip: # we already rotated images for 270' for b in range(len(images)): images_global[b] = transforms.functional.rotate( images_global[b], 90 ) # rotate back! images_global[b] = transforms.functional.hflip(images_global[b]) if self.mode == 2 or self.mode == 3: images_local[b] = transforms.functional.rotate( images_local[b], 90 ) # rotate back! images_local[b] = transforms.functional.hflip( images_local[b] ) for angle in self.rotate_range: if angle > 0: for b in range(len(images)): images_global[b] = transforms.functional.rotate( images_global[b], 90 ) if self.mode == 2 or self.mode == 3: images_local[b] = transforms.functional.rotate( images_local[b], 90 ) # prepare global images onto cuda images_glb = images_transform(images_global) # b, c, h, w if self.mode == 2 or self.mode == 3: # patches = global2patch(images_local, self.n, self.step, self.size_p) patches, coordinates, templates, sizes, ratios = global2patch( # patches, coordinates, _, sizes, ratios = global2patch( images, self.size_p ) # predicted_patches = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] # predicted_ensembles = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] predicted_patches = [ np.zeros( ( len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1], ) ) for i in range(len(images)) ] predicted_ensembles = [ np.zeros( ( len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1], ) ) for i in range(len(images)) ] if self.mode == 1: # eval with only resized global image if flip: outputs_global += np.flip( np.rot90( model.forward(images_glb, None, None, None)[0] .data.cpu() .numpy(), k=angle, axes=(3, 2), ), axis=3, ) else: outputs_global += np.rot90( model.forward(images_glb, None, None, None)[0] .data.cpu() .numpy(), k=angle, axes=(3, 2), ) if self.mode == 2: # eval with patches for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): patches_var = images_transform( patches[i][j : j + self.sub_batch_size] ) # b, c, h, w # output_ensembles, output_global, output_patches, _ = model.forward(images_glb[i:i+1], patches_var, self.coordinates[j : j+self.sub_batch_size], self.ratio, mode=self.mode, n_patch=self.n**2) # include cordinates output_ensembles, output_global, output_patches, _ = model.forward( images_glb[i : i + 1], patches_var, coordinates[i][j : j + self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i]), ) # patch predictions predicted_patches[i][ j : j + output_patches.size()[0] ] += ( F.interpolate( output_patches, size=self.size_p, mode="nearest" ) .data.cpu() .numpy() ) predicted_ensembles[i][ j : j + output_ensembles.size()[0] ] += ( F.interpolate( output_ensembles, size=self.size_p, mode="nearest", ) .data.cpu() .numpy() ) j += patches_var.size()[0] if flip: outputs_global[i] += np.flip( np.rot90( output_global[0].data.cpu().numpy(), k=angle, axes=(2, 1), ), axis=2, ) # scores_local[i] += np.flip(np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps) # scores[i] += np.flip(np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps) scores_local[i] += np.flip( np.rot90( np.array( patch2global( predicted_patches[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ), axis=3, ) # merge softmax scores from patches (overlaps) scores[i] += np.flip( np.rot90( np.array( patch2global( predicted_ensembles[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ), axis=3, ) # merge softmax scores from patches (overlaps) else: outputs_global[i] += np.rot90( output_global[0].data.cpu().numpy(), k=angle, axes=(2, 1), ) # scores_local[i] += np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps) # scores[i] += np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps) scores_local[i] += np.rot90( np.array( patch2global( predicted_patches[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ) # merge softmax scores from patches (overlaps) scores[i] += np.rot90( np.array( patch2global( predicted_ensembles[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ) # merge softmax scores from patches (overlaps) if self.mode == 3: # eval global with help from patches # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): patches_var = images_transform( patches[i][j : j + self.sub_batch_size] ) # b, c, h, w # fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates fm_patches, output_patches = model.module.collect_local_fm( images_glb[i : i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, template=templates[0], n_patch_all=len(coordinates[i]), ) predicted_patches[i][ j : j + output_patches.size()[0] ] += ( F.interpolate( output_patches, size=self.size_p, mode="nearest" ) .data.cpu() .numpy() ) j += self.sub_batch_size # go through global image # tmp, fm_global = model.forward(images_glb, None, self.coordinates, self.ratio, mode=self.mode, global_model=None, n_patch=self.n**2) # include cordinates tmp, fm_global = model.forward( images_glb, None, None, None, mode=self.mode ) if flip: outputs_global += np.flip( np.rot90(tmp.data.cpu().numpy(), k=angle, axes=(3, 2)), axis=3, ) else: outputs_global += np.rot90( tmp.data.cpu().numpy(), k=angle, axes=(3, 2) ) # generate ensembles for i in range(len(images)): j = 0 # while j < self.n ** 2: while j < len(coordinates[i]): fl = fm_patches[i][j : j + self.sub_batch_size].cuda() # fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0] fg = model.module._crop_global( fm_global[i : i + 1], coordinates[i][j : j + self.sub_batch_size], ratios[i], )[0] fg = F.interpolate( fg, size=fl.size()[2:], mode="bilinear" ) output_ensembles = model.module.ensemble( fl, fg ) # include cordinates # output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs) # ensemble predictions predicted_ensembles[i][ j : j + output_ensembles.size()[0] ] += ( F.interpolate( output_ensembles, size=self.size_p, mode="nearest", ) .data.cpu() .numpy() ) j += self.sub_batch_size if flip: # scores_local[i] += np.flip(np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps) # scores[i] += np.flip(np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps) scores_local[i] += np.flip( np.rot90( np.array( patch2global( predicted_patches[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ), axis=3, )[ 0 ] # merge softmax scores from patches (overlaps) scores[i] += np.flip( np.rot90( np.array( patch2global( predicted_ensembles[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ), axis=3, )[ 0 ] # merge softmax scores from patches (overlaps) else: # scores_local[i] += np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps) # scores[i] += np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps) scores_local[i] += np.rot90( np.array( patch2global( predicted_patches[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ) # merge softmax scores from patches (overlaps) scores[i] += np.rot90( np.array( patch2global( predicted_ensembles[i : i + 1], self.n_class, sizes[i : i + 1], coordinates[i : i + 1], self.size_p, ) ), k=angle, axes=(3, 2), ) # merge softmax scores from patches (overlaps) # global predictions # predictions_global = F.interpolate(torch.Tensor(outputs_global), self.size0, mode='nearest').argmax(1).detach().numpy() outputs_global = torch.Tensor(outputs_global) predictions_global = [ F.interpolate( outputs_global[i : i + 1], images[i].size[::-1], mode="nearest" ) .argmax(1) .detach() .numpy()[0] for i in range(len(images)) ] if not self.test: self.metrics_global.update(labels_npy, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions # predictions_local = scores_local.argmax(1) # b, h, w predictions_local = [score.argmax(1)[0] for score in scores_local] if not self.test: self.metrics_local.update(labels_npy, predictions_local) # combined/ensemble predictions # predictions = scores.argmax(1) # b, h, w predictions = [score.argmax(1)[0] for score in scores] if not self.test: self.metrics.update(labels_npy, predictions) return predictions, predictions_global, predictions_local else: return None, predictions_global, None
class Trainer(object): # def __init__(self, criterion, optimizer, n_class, size0, size_g, size_p, n, sub_batch_size=6, mode=1, lamb_fmreg=0.15): def __init__( self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15, ): self.criterion = criterion self.optimizer = optimizer self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class # self.size0 = size0 self.size_g = size_g self.size_p = size_p # self.n = n self.sub_batch_size = sub_batch_size self.mode = mode self.lamb_fmreg = lamb_fmreg # self.ratio = float(size_p[0]) / size0[0] # self.step = (size0[0] - size_p[0]) // (n - 1) # self.template, self.coordinates = template_patch2global(size0, size_p, n, self.step) def set_train(self, model): model.module.ensemble_conv.train() if self.mode == 1 or self.mode == 3: model.module.resnet_global.train() model.module.fpn_global.train() else: model.module.resnet_local.train() model.module.fpn_local.train() def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def train(self, sample, model, global_fixed): images, labels = sample["image"], sample["label"] # PIL images #lbls = [RGB_mapping_to_class(np.array(label)) for label in labels] #labels = [Image.fromarray(lbl) for lbl in lbls] # del(lbls) labels_npy = masks_transform( labels, numpy=True ) # label of origin size in numpy images_glb = resize(images, self.size_g) # list of resized PIL images images_glb = images_transform(images_glb) labels_glb = resize( labels, (self.size_g[0] // 4, self.size_g[1] // 4), label=True ) # down 1/4 for loss # labels_glb = resize(labels, self.size_g, label=True) # must downsample image for reduced GPU memory labels_glb = masks_transform(labels_glb) # 127 * 127 * 8 = 129032 if self.mode == 2 or self.mode == 3: patches, coordinates, templates, sizes, ratios = global2patch( # patches, coordinates, _, sizes, ratios = global2patch( images, self.size_p ) label_patches, _, _, _, _ = global2patch(labels, self.size_p) # patches, label_patches = global2patch(images, self.n, self.step, self.size_p), global2patch(labels, self.n, self.step, self.size_p) # predicted_patches = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] # predicted_ensembles = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] predicted_patches = [ np.zeros( (len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1]) ) for i in range(len(images)) ] predicted_ensembles = [ np.zeros( (len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1]) ) for i in range(len(images)) ] outputs_global = [None for i in range(len(images))] if self.mode == 1: # training with only (resized) global image outputs_global, _ = model.forward(images_glb, None, None, None) loss = self.criterion(outputs_global, labels_glb) loss.backward() self.optimizer.step() self.optimizer.zero_grad() if self.mode == 2: # training with patches for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): patches_var = images_transform( patches[i][j : j + self.sub_batch_size] ) # b, c, h, w label_patches_var = masks_transform( resize( label_patches[i][j : j + self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True, ) ) # down 1/4 for loss # label_patches_var = masks_transform(label_patches[i][j : j+self.sub_batch_size]) # output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, self.coordinates[j : j+self.sub_batch_size], self.ratio, mode=self.mode, n_patch=self.n**2) # include cordinates output_ensembles, output_global, output_patches, fmreg_l2 = model.forward( images_glb[i : i + 1], patches_var, coordinates[i][j : j + self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i]), ) loss = ( self.criterion(output_patches, label_patches_var) + self.criterion(output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2 ) loss.backward() # patch predictions predicted_patches[i][j : j + output_patches.size()[0]] = ( F.interpolate(output_patches, size=self.size_p, mode="nearest") .data.cpu() .numpy() ) predicted_ensembles[i][j : j + output_ensembles.size()[0]] = ( F.interpolate( output_ensembles, size=self.size_p, mode="nearest" ) .data.cpu() .numpy() ) j += self.sub_batch_size outputs_global[i] = output_global outputs_global = torch.cat(outputs_global, dim=0) self.optimizer.step() self.optimizer.zero_grad() if self.mode == 3: # train global with help from patches # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): patches_var = images_transform( patches[i][j : j + self.sub_batch_size] ) # b, c, h, w # fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates fm_patches, output_patches = model.module.collect_local_fm( images_glb[i : i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, template=templates[0], n_patch_all=len(coordinates[i]), ) predicted_patches[i][j : j + output_patches.size()[0]] = ( F.interpolate(output_patches, size=self.size_p, mode="nearest") .data.cpu() .numpy() ) j += self.sub_batch_size # train on global image outputs_global, fm_global = model.forward( images_glb, None, None, None, mode=self.mode ) loss = self.criterion(outputs_global, labels_glb) loss.backward(retain_graph=True) # fmreg loss # generate ensembles & calc loss for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): label_patches_var = masks_transform( resize( label_patches[i][j : j + self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True, ) ) # label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], self.size_p, label=True)) fl = fm_patches[i][j : j + self.sub_batch_size].cuda() # fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0] fg = model.module._crop_global( fm_global[i : i + 1], coordinates[i][j : j + self.sub_batch_size], ratios[i], )[0] fg = F.interpolate(fg, size=fl.size()[2:], mode="bilinear") output_ensembles = model.module.ensemble(fl, fg) # output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs) loss = self.criterion( output_ensembles, label_patches_var ) # + 0.15 * mse(fl, fg) # if i == len(images) - 1 and j + self.sub_batch_size >= self.n**2: if i == len(images) - 1 and j + self.sub_batch_size >= len( coordinates[i] ): loss.backward() else: loss.backward(retain_graph=True) # ensemble predictions predicted_ensembles[i][j : j + output_ensembles.size()[0]] = ( F.interpolate( output_ensembles, size=self.size_p, mode="nearest" ) .data.cpu() .numpy() ) j += self.sub_batch_size self.optimizer.step() self.optimizer.zero_grad() # global predictions # predictions_global = F.interpolate(outputs_global.cpu(), self.size0, mode='nearest').argmax(1).detach().numpy() outputs_global = outputs_global.cpu() predictions_global = [ F.interpolate( outputs_global[i : i + 1], images[i].size[::-1], mode="nearest" ) .argmax(1) .detach() .numpy() for i in range(len(images)) ] self.metrics_global.update(labels_npy, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions # scores_local = np.array(patch2global(predicted_patches, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps) scores_local = np.array( patch2global( predicted_patches, self.n_class, sizes, coordinates, self.size_p ) ) # merge softmax scores from patches (overlaps) predictions_local = scores_local.argmax(1) # b, h, w self.metrics_local.update(labels_npy, predictions_local) # combined/ensemble predictions # scores = np.array(patch2global(predicted_ensembles, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps) scores = np.array( patch2global( predicted_ensembles, self.n_class, sizes, coordinates, self.size_p ) ) # merge softmax scores from patches (overlaps) predictions = scores.argmax(1) # b, h, w self.metrics.update(labels_npy, predictions) return loss
class Evaluator(object): def __init__(self, n_class, size_g, size_p, sub_batch_size=6, mode=1, test=False): self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class self.size_g = size_g self.size_p = size_p self.sub_batch_size = sub_batch_size self.mode = mode self.test = test if test: self.flip_range = [False, True] self.rotate_range = [0, 1, 2, 3] else: self.flip_range = [False] self.rotate_range = [0] def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def eval_test(self, sample, model, global_fixed): with torch.no_grad(): images = sample['image'] if not self.test: labels = sample['label'] # PIL images labels_npy = masks_transform(labels, numpy=True) images_global = resize(images, self.size_g) outputs_global = np.zeros( (len(images), self.n_class, self.size_g[0] // 4, self.size_g[1] // 4)) if self.mode == 2 or self.mode == 3: images_local = [image.copy() for image in images] scores_local = [ np.zeros((1, self.n_class, images[i].size[1], images[i].size[0])) for i in range(len(images)) ] scores = [ np.zeros((1, self.n_class, images[i].size[1], images[i].size[0])) for i in range(len(images)) ] for flip in self.flip_range: if flip: # we already rotated images for 270' for b in range(len(images)): images_global[b] = transforms.functional.rotate( images_global[b], 90) # rotate back! images_global[b] = transforms.functional.hflip( images_global[b]) if self.mode == 2 or self.mode == 3: images_local[b] = transforms.functional.rotate( images_local[b], 90) # rotate back! images_local[b] = transforms.functional.hflip( images_local[b]) for angle in self.rotate_range: if angle > 0: for b in range(len(images)): images_global[b] = transforms.functional.rotate( images_global[b], 90) if self.mode == 2 or self.mode == 3: images_local[b] = transforms.functional.rotate( images_local[b], 90) # prepare global images onto cuda images_glb = images_transform(images_global) # b, c, h, w if self.mode == 2 or self.mode == 3: patches, coordinates, templates, sizes, ratios = global2patch( images, self.size_p) predicted_patches = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] predicted_ensembles = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] if self.mode == 1: # eval with only resized global image ########################## if flip: outputs_global += np.flip(np.rot90(model.forward( images_glb, None, None, None)[0].data.cpu().numpy(), k=angle, axes=(3, 2)), axis=3) else: outputs_global += np.rot90(model.forward( images_glb, None, None, None)[0].data.cpu().numpy(), k=angle, axes=(3, 2)) ################################################################ if self.mode == 2: # eval with patches ########################################### for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i] [j:j + self.sub_batch_size]) # b, c, h, w output_ensembles, output_global, output_patches, _ = model.forward( images_glb[i:i + 1], patches_var, coordinates[i][j:j + self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i])) # patch predictions predicted_patches[i][j:j + output_patches.size( )[0]] += F.interpolate( output_patches, size=self.size_p, mode='nearest').data.cpu().numpy() predicted_ensembles[i][j:j + output_ensembles.size( )[0]] += F.interpolate( output_ensembles, size=self.size_p, mode='nearest' ).data.cpu().numpy() j += patches_var.size()[0] if flip: outputs_global[i] += np.flip(np.rot90( output_global[0].data.cpu().numpy(), k=angle, axes=(2, 1)), axis=2) scores_local[i] += np.flip( np.rot90(np.array( patch2global( predicted_patches[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2)), axis=3 ) # merge softmax scores from patches (overlaps) scores[i] += np.flip( np.rot90(np.array( patch2global( predicted_ensembles[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2)), axis=3 ) # merge softmax scores from patches (overlaps) else: outputs_global[i] += np.rot90( output_global[0].data.cpu().numpy(), k=angle, axes=(2, 1)) scores_local[i] += np.rot90( np.array( patch2global( predicted_patches[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2) ) # merge softmax scores from patches (overlaps) scores[i] += np.rot90( np.array( patch2global( predicted_ensembles[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2) ) # merge softmax scores from patches (overlaps) ############################################################### if self.mode == 3: # eval global with help from patches ################################################## # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i] [j:j + self.sub_batch_size]) # b, c, h, w fm_patches, output_patches = model.module.collect_local_fm( images_glb[i:i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, template=templates[i], n_patch_all=len(coordinates[i])) predicted_patches[i][j:j + output_patches.size( )[0]] += F.interpolate( output_patches, size=self.size_p, mode='nearest').data.cpu().numpy() j += self.sub_batch_size # go through global image tmp, fm_global = model.forward(images_glb, None, None, None, mode=self.mode) if flip: outputs_global += np.flip(np.rot90( tmp.data.cpu().numpy(), k=angle, axes=(3, 2)), axis=3) else: outputs_global += np.rot90(tmp.data.cpu().numpy(), k=angle, axes=(3, 2)) # generate ensembles for i in range(len(images)): j = 0 while j < len(coordinates[i]): fl = fm_patches[i][j:j + self.sub_batch_size].cuda() fg = model.module._crop_global( fm_global[i:i + 1], coordinates[i][j:j + self.sub_batch_size], ratios[i])[0] fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear') output_ensembles = model.module.ensemble( fl, fg) # include cordinates # ensemble predictions predicted_ensembles[i][j:j + output_ensembles.size( )[0]] += F.interpolate( output_ensembles, size=self.size_p, mode='nearest' ).data.cpu().numpy() j += self.sub_batch_size if flip: scores_local[i] += np.flip( np.rot90(np.array( patch2global( predicted_patches[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2)), axis=3 )[0] # merge softmax scores from patches (overlaps) scores[i] += np.flip( np.rot90(np.array( patch2global( predicted_ensembles[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2)), axis=3 )[0] # merge softmax scores from patches (overlaps) else: scores_local[i] += np.rot90( np.array( patch2global( predicted_patches[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2) ) # merge softmax scores from patches (overlaps) scores[i] += np.rot90( np.array( patch2global( predicted_ensembles[i:i + 1], self.n_class, sizes[i:i + 1], coordinates[i:i + 1], self.size_p)), k=angle, axes=(3, 2) ) # merge softmax scores from patches (overlaps) ################################################### # global predictions ########################### outputs_global = torch.Tensor(outputs_global) predictions_global = [ F.interpolate(outputs_global[i:i + 1], images[i].size[::-1], mode='nearest').argmax(1).detach().numpy()[0] for i in range(len(images)) ] if not self.test: self.metrics_global.update(labels_npy, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions ########################### predictions_local = [ score.argmax(1)[0] for score in scores_local ] if not self.test: self.metrics_local.update(labels_npy, predictions_local) ################################################### # combined/ensemble predictions ########################### predictions = [score.argmax(1)[0] for score in scores] if not self.test: self.metrics.update(labels_npy, predictions) return predictions, predictions_global, predictions_local else: return None, predictions_global, None
class Trainer(object): def __init__(self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15): self.criterion = criterion self.optimizer = optimizer self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class self.size_g = size_g self.size_p = size_p self.sub_batch_size = sub_batch_size self.mode = mode self.lamb_fmreg = lamb_fmreg def set_train(self, model): model.module.ensemble_conv.train() if self.mode == 1 or self.mode == 3: model.module.resnet_global.train() model.module.fpn_global.train() else: model.module.resnet_local.train() model.module.fpn_local.train() def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def train(self, sample, model, global_fixed): images, labels = sample['image'], sample['label'] # PIL images labels_npy = masks_transform( labels, numpy=True) # label of origin size in numpy images_glb = resize(images, self.size_g) # list of resized PIL images images_glb = images_transform(images_glb) labels_glb = resize(labels, (self.size_g[0] // 4, self.size_g[1] // 4), label=True) # FPN down 1/4, for loss labels_glb = masks_transform(labels_glb) if self.mode == 2 or self.mode == 3: patches, coordinates, templates, sizes, ratios = global2patch( images, self.size_p) label_patches, _, _, _, _ = global2patch(labels, self.size_p) predicted_patches = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] predicted_ensembles = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] outputs_global = [None for i in range(len(images))] if self.mode == 1: # training with only (resized) global image ######################################### outputs_global, _ = model.forward(images_glb, None, None, None) loss = self.criterion(outputs_global, labels_glb) loss.backward() self.optimizer.step() self.optimizer.zero_grad() ############################################## if self.mode == 2: # training with patches ########################################### for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i][j:j + self.sub_batch_size]) # b, c, h, w label_patches_var = masks_transform( resize(label_patches[i][j:j + self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True)) # down 1/4 for loss output_ensembles, output_global, output_patches, fmreg_l2 = model.forward( images_glb[i:i + 1], patches_var, coordinates[i][j:j + self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i])) loss = self.criterion( output_patches, label_patches_var) + self.criterion( output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2 loss.backward() # patch predictions predicted_patches[i][j:j + output_patches.size( )[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy() predicted_ensembles[i][j:j + output_ensembles.size( )[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy() j += self.sub_batch_size outputs_global[i] = output_global outputs_global = torch.cat(outputs_global, dim=0) self.optimizer.step() self.optimizer.zero_grad() ##################################################################################### if self.mode == 3: # train global with help from patches ################################################## # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i][j:j + self.sub_batch_size]) # b, c, h, w fm_patches, output_patches = model.module.collect_local_fm( images_glb[i:i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, template=templates[i], n_patch_all=len(coordinates[i])) predicted_patches[i][j:j + output_patches.size( )[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy() j += self.sub_batch_size # train on global image outputs_global, fm_global = model.forward(images_glb, None, None, None, mode=self.mode) loss = self.criterion(outputs_global, labels_glb) loss.backward(retain_graph=True) # fmreg loss # generate ensembles & calc loss for i in range(len(images)): j = 0 while j < len(coordinates[i]): label_patches_var = masks_transform( resize(label_patches[i][j:j + self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True)) fl = fm_patches[i][j:j + self.sub_batch_size].cuda() fg = model.module._crop_global( fm_global[i:i + 1], coordinates[i][j:j + self.sub_batch_size], ratios[i])[0] fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear') output_ensembles = model.module.ensemble(fl, fg) loss = self.criterion( output_ensembles, label_patches_var) # + 0.15 * mse(fl, fg) if i == len(images) - 1 and j + self.sub_batch_size >= len( coordinates[i]): loss.backward() else: loss.backward(retain_graph=True) # ensemble predictions predicted_ensembles[i][j:j + output_ensembles.size( )[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy() j += self.sub_batch_size self.optimizer.step() self.optimizer.zero_grad() # global predictions ########################### outputs_global = outputs_global.cpu() predictions_global = [ F.interpolate(outputs_global[i:i + 1], images[i].size[::-1], mode='nearest').argmax(1).detach().numpy() for i in range(len(images)) ] self.metrics_global.update(labels_npy, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions ########################### scores_local = np.array( patch2global(predicted_patches, self.n_class, sizes, coordinates, self.size_p) ) # merge softmax scores from patches (overlaps) predictions_local = scores_local.argmax(1) # b, h, w self.metrics_local.update(labels_npy, predictions_local) ################################################### # combined/ensemble predictions ########################### scores = np.array( patch2global(predicted_ensembles, self.n_class, sizes, coordinates, self.size_p) ) # merge softmax scores from patches (overlaps) predictions = scores.argmax(1) # b, h, w self.metrics.update(labels_npy, predictions) return loss
class Trainer(object): # def __init__(self, criterion, optimizer, n_class, size0, size_g, size_p, n, sub_batch_size=6, mode=1, lamb_fmreg=0.15): def __init__(self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15): self.criterion = criterion self.optimizer = optimizer self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class # self.size0 = size0 self.size_g = size_g self.size_p = size_p # self.n = n self.sub_batch_size = sub_batch_size self.mode = mode self.lamb_fmreg = lamb_fmreg # self.ratio = float(size_p[0]) / size0[0] # self.step = (size0[0] - size_p[0]) // (n - 1) # self.template, self.coordinates = template_patch2global(size0, size_p, n, self.step) def set_train(self, model, parallel=True): if not parallel: model.ensemble_conv.train() if self.mode == 1 or self.mode == 3: model.resnet_global.train() model.fpn_global.train() else: model.resnet_local.train() model.fpn_local.train() else: model.module.ensemble_conv.train() if self.mode == 1 or self.mode == 3: model.module.resnet_global.train() model.module.fpn_global.train() else: model.module.resnet_local.train() model.module.fpn_local.train() def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def train(self, sample, model, global_fixed, parallel=True): images, labels = sample['image'], sample['label'] # PIL images labels_npy = masks_transform(labels, numpy=True) # label of origin size in numpy images_glb = resize(images, self.size_g) # list of resized PIL images images_glb = images_transform(images_glb) labels_glb = resize(labels, (self.size_g[0] // 4, self.size_g[1] // 4), label=True) # down 1/4 for loss # labels_glb = resize(labels, self.size_g, label=True) # must downsample image for reduced GPU memory labels_glb = masks_transform(labels_glb) if self.mode == 2 or self.mode == 3: patches, coordinates, templates, sizes, ratios = global2patch(images, self.size_p) label_patches, _, _, _, _ = global2patch(labels, self.size_p) # patches, label_patches = global2patch(images, self.n, self.step, self.size_p), global2patch(labels, self.n, self.step, self.size_p) # predicted_patches = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] # predicted_ensembles = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] predicted_patches = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] predicted_ensembles = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ] outputs_global = [ None for i in range(len(images)) ] if self.mode == 1: # training with only (resized) global image ######################################### outputs_global, _ = model.forward(images_glb, None, None, None) loss = self.criterion(outputs_global, labels_glb) loss.backward() self.optimizer.step() self.optimizer.zero_grad() ############################################## if self.mode == 2: # training with patches ########################################### for i in range(len(images)): # sync the start for each global image torch.distributed.barrier() #print("Will start training global image:", i, " on rank :", torch.distributed.get_rank()) #group = torch.distributed.new_group(range(torch.distributed.get_world_size())) coordinate_size = torch.tensor([len(coordinates[i])], dtype=torch.int32).cuda() torch.distributed.all_reduce(coordinate_size, op=torch.distributed.ReduceOp.MAX) #print("Will train:", coordinate_size.item(), "loops in global image.", "On rank :", torch.distributed.get_rank()) j = 0 _loop = 0 # while j < self.n**2: #print("==== ==== len(coordinates[i]):", len(coordinates[i]), ", at rank:", torch.distributed.get_rank()) #while j < len(coordinates[i]): #while j in range(6): #while _ in range(coordinate_size.item()): while _loop < coordinate_size.item(): patches_var = images_transform(patches[i][j : j+self.sub_batch_size]) # b, c, h, w label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True)) # down 1/4 for loss # label_patches_var = masks_transform(label_patches[i][j : j+self.sub_batch_size]) # output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, self.coordinates[j : j+self.sub_batch_size], self.ratio, mode=self.mode, n_patch=self.n**2) # include cordinates output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, coordinates[i][j : j+self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i])) loss = self.criterion(output_patches, label_patches_var) + self.criterion(output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2 loss.backward() # patch predictions predicted_patches[i][j:j+output_patches.size()[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy() predicted_ensembles[i][j:j+output_ensembles.size()[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy() # Because we choose loop the biggest coordinate_size in all ranks, # make sure not cross the border for current rank j = (j+self.sub_batch_size)%len(coordinates[i]) _loop += self.sub_batch_size #print("==== ==== nested loop:", j, ", at rank:", torch.distributed.get_rank()) outputs_global[i] = output_global outputs_global = torch.cat(outputs_global, dim=0) self.optimizer.step() self.optimizer.zero_grad() ##################################################################################### if self.mode == 3: # train global with help from patches ################################################## # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): patches_var = images_transform(patches[i][j : j+self.sub_batch_size]) # b, c, h, w if parallel: # fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, ratios[i], coordinates[i], [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=len(coordinates[i])) else: # fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates fm_patches, output_patches = model.collect_local_fm(images_glb[i:i+1], patches_var, ratios[i], coordinates[i], [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=len(coordinates[i])) predicted_patches[i][j:j+output_patches.size()[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy() j += self.sub_batch_size # train on global image outputs_global, fm_global = model.forward(images_glb, None, None, None, mode=self.mode) loss = self.criterion(outputs_global, labels_glb) loss.backward(retain_graph=True) # fmreg loss # generate ensembles & calc loss for i in range(len(images)): j = 0 # while j < self.n**2: while j < len(coordinates[i]): label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True)) # label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], self.size_p, label=True)) fl = fm_patches[i][j : j+self.sub_batch_size].cuda() if parallel: # fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0] fg = model.module._crop_global(fm_global[i:i+1], coordinates[i][j:j+self.sub_batch_size], ratios[i])[0] else: # fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0] fg = model._crop_global(fm_global[i:i+1], coordinates[i][j:j+self.sub_batch_size], ratios[i])[0] fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear') if parallel: output_ensembles = model.module.ensemble(fl, fg) # output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs) else: output_ensembles = model.ensemble(fl, fg) # output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs) loss = self.criterion(output_ensembles, label_patches_var)# + 0.15 * mse(fl, fg) # if i == len(images) - 1 and j + self.sub_batch_size >= self.n**2: if i == len(images) - 1 and j + self.sub_batch_size >= len(coordinates[i]): loss.backward() else: loss.backward(retain_graph=True) # ensemble predictions predicted_ensembles[i][j:j+output_ensembles.size()[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy() j += self.sub_batch_size self.optimizer.step() self.optimizer.zero_grad() # global predictions ########################### # predictions_global = F.interpolate(outputs_global.cpu(), self.size0, mode='nearest').argmax(1).detach().numpy() outputs_global = outputs_global.cpu() predictions_global = [F.interpolate(outputs_global[i:i+1], images[i].size[::-1], mode='nearest').argmax(1).detach().numpy() for i in range(len(images))] self.metrics_global.update(labels_npy, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions ########################### # scores_local = np.array(patch2global(predicted_patches, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps) scores_local = np.array(patch2global(predicted_patches, self.n_class, sizes, coordinates, self.size_p)) # merge softmax scores from patches (overlaps) predictions_local = scores_local.argmax(1) # b, h, w self.metrics_local.update(labels_npy, predictions_local) ################################################### # combined/ensemble predictions ########################### # scores = np.array(patch2global(predicted_ensembles, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps) scores = np.array(patch2global(predicted_ensembles, self.n_class, sizes, coordinates, self.size_p)) # merge softmax scores from patches (overlaps) predictions = scores.argmax(1) # b, h, w self.metrics.update(labels_npy, predictions) return loss
class Trainer(object): def __init__(self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15): self.criterion = criterion self.optimizer = optimizer self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class self.size_g = size_g self.size_p = size_p self.sub_batch_size = sub_batch_size self.mode = mode self.lamb_fmreg = lamb_fmreg def set_train(self, model): model.module.ensemble_conv.train() if self.mode == 1 or self.mode == 3: model.module.resnet_global.train() model.module.fpn_global.train() else: model.module.resnet_local.train() model.module.fpn_local.train() def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def train(self, sample, model, global_fixed): images, labels = sample['image'], sample['label'] # PIL images ids = sample['id'] width, height = images[0].size if width != self.size_g[0] or height != self.size_g[1]: images_glb = resize(images, self.size_g) # list of resized PIL images else: images_glb = list(images) images_glb = images_transform(images_glb) labels_glb = masks_transform(labels) if self.mode == 2 or self.mode == 3: patches, coordinates, sizes, ratios, label_patches = global2patch( images, labels_glb, self.size_p, ids) predicted_patches = np.zeros((len(images), 4)) predicted_ensembles = np.zeros((len(images), 4)) outputs_global = [None for i in range(len(images))] if self.mode == 1: # training with only (resized) global image ######################################### outputs_global, _ = model.forward(images_glb, None, None, None) loss = self.criterion(outputs_global, labels_glb) loss.backward() self.optimizer.step() self.optimizer.zero_grad() ############################################## if self.mode == 2: # training with patches ########################################### for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i][j:j + self.sub_batch_size]) # b, c, h, w label_patches_var = masks_transform( label_patches[i][j:j + self.sub_batch_size]) output_ensembles, output_global, output_patches, fmreg_l2 = model.forward( images_glb[i:i + 1], patches_var, coordinates[i][j:j + self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i])) loss = self.criterion( output_patches, label_patches_var) + self.criterion( output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2 loss.backward() # patch predictions for pred_patch, pred_ensemble in zip( torch.max(output_patches.data, 1)[1].data, torch.max(output_ensembles.data, 1)[1].data): predicted_patches[i][int(pred_patch)] += 1 predicted_ensembles[i][int(pred_ensemble)] += 1 j += self.sub_batch_size outputs_global[i] = output_global outputs_global = torch.cat(outputs_global, dim=0) self.optimizer.step() self.optimizer.zero_grad() ##################################################################################### if self.mode == 3: # train global with help from patches ################################################## # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i][j:j + self.sub_batch_size]) # b, c, h, w _, output_patches = model.module.collect_local_fm( images_glb[i:i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, n_patch_all=len(coordinates[i])) for pred_patch in torch.max(output_patches.data, 1)[1].data: predicted_patches[i][int(pred_patch)] += 1 j += self.sub_batch_size # train on global image outputs_global, fm_global = model.forward(images_glb, None, None, None, mode=self.mode) loss = self.criterion(outputs_global, labels_glb) loss.backward(retain_graph=True) # fmreg loss # generate ensembles & calc loss for i in range(len(images)): j = 0 while j < len(coordinates[i]): label_patches_var = masks_transform( label_patches[i][j:j + self.sub_batch_size]) patches_var = images_transform( patches[i][j:j + self.sub_batch_size]) # b, c, h, w fl = model.module.generate_local_fm( images_glb[i:i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, n_patch_all=len(coordinates[i])) fg = model.module._crop_global( fm_global[i:i + 1], coordinates[i][j:j + self.sub_batch_size], ratios[i])[0] fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear') output_ensembles = model.module.ensemble(fl, fg) loss = self.criterion( output_ensembles, label_patches_var) # + 0.15 * mse(fl, fg) if i == len(images) - 1 and j + self.sub_batch_size >= len( coordinates[i]): loss.backward() else: loss.backward(retain_graph=True) # ensemble predictions for pred_ensemble in torch.max(output_ensembles.data, 1)[1].data: predicted_ensembles[i][int(pred_ensemble)] += 1 j += self.sub_batch_size self.optimizer.step() self.optimizer.zero_grad() # global predictions ########################### _, predictions_global = torch.max(outputs_global.data, 1) self.metrics_global.update(labels_glb, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions ########################### predictions_local = predicted_patches.argmax(1) #self.metrics_local.update(labels_npy, predictions_local) self.metrics_local.update(labels_glb, predictions_local) ################################################### # combined/ensemble predictions ########################### predictions = predicted_ensembles.argmax(1) self.metrics.update(labels_glb, predictions) return loss
class Evaluator(object): def __init__(self, n_class, size_g, size_p, sub_batch_size=6, mode=1, test=False): self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class self.size_g = size_g self.size_p = size_p self.sub_batch_size = sub_batch_size self.mode = mode self.test = test if test: self.flip_range = [False, True] self.rotate_range = [0, 1, 2, 3] else: self.flip_range = [False] self.rotate_range = [0] def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def eval_test(self, sample, model, global_fixed): with torch.no_grad(): images = sample['image'] ids = sample['id'] if not self.test: labels = sample['label'] # PIL images labels_glb = masks_transform(labels) width, height = images[0].size if width > self.size_g[0] or height > self.size_g[1]: images_global = resize( images, self.size_g) # list of resized PIL images else: images_global = list(images) if self.mode == 2 or self.mode == 3: images_local = [image.copy() for image in images] scores_local = [ np.zeros((1, self.n_class, images[i].size[1], images[i].size[0])) for i in range(len(images)) ] scores = [ np.zeros((1, self.n_class, images[i].size[1], images[i].size[0])) for i in range(len(images)) ] for flip in self.flip_range: if flip: # we already rotated images for 270' for b in range(len(images)): images_global[b] = transforms.functional.rotate( images_global[b], 90) # rotate back! images_global[b] = transforms.functional.hflip( images_global[b]) if self.mode == 2 or self.mode == 3: images_local[b] = transforms.functional.rotate( images_local[b], 90) # rotate back! images_local[b] = transforms.functional.hflip( images_local[b]) for angle in self.rotate_range: if angle > 0: for b in range(len(images)): images_global[b] = transforms.functional.rotate( images_global[b], 90) if self.mode == 2 or self.mode == 3: images_local[b] = transforms.functional.rotate( images_local[b], 90) # prepare global images onto cuda images_glb = images_transform(images_global) # b, c, h, w if self.mode == 2 or self.mode == 3: patches, coordinates, sizes, ratios, label_patches = global2patch( images, labels_glb, self.size_p, ids) predicted_patches = np.zeros((len(images), 4)) predicted_ensembles = np.zeros((len(images), 4)) outputs_global = [None for i in range(len(images))] if self.mode == 1: # eval with only resized global image ########################## if flip: outputs_global += np.flip(np.rot90(model.forward( images_glb, None, None, None)[0].data.cpu().numpy(), k=angle, axes=(3, 2)), axis=3) else: outputs_global, _ = model.forward( images_glb, None, None, None) ################################################################ if self.mode == 2: # eval with patches ########################################### for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i] [j:j + self.sub_batch_size]) # b, c, h, w output_ensembles, output_global, output_patches, _ = model.forward( images_glb[i:i + 1], patches_var, coordinates[i][j:j + self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i])) # patch predictions for pred_patch, pred_ensemble in zip( torch.max(output_patches.data, 1)[1].data, torch.max(output_ensembles.data, 1)[1].data): predicted_patches[i][int(pred_patch)] += 1 predicted_ensembles[i][int( pred_ensemble)] += 1 j += patches_var.size()[0] outputs_global[i] = output_global outputs_global = torch.cat(outputs_global, dim=0) ############################################################### if self.mode == 3: # eval global with help from patches ################################################## # go through local patches to collect feature maps # collect predictions from patches for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i] [j:j + self.sub_batch_size]) # b, c, h, w _, output_patches = model.module.collect_local_fm( images_glb[i:i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, n_patch_all=len(coordinates[i])) for pred_patch in torch.max( output_patches.data, 1)[1].data: predicted_patches[i][int(pred_patch)] += 1 j += self.sub_batch_size # go through global image tmp, fm_global = model.forward(images_glb, None, None, None, mode=self.mode) if flip: outputs_global += np.flip(np.rot90( tmp.data.cpu().numpy(), k=angle, axes=(3, 2)), axis=3) else: outputs_global = tmp # generate ensembles for i in range(len(images)): j = 0 while j < len(coordinates[i]): patches_var = images_transform( patches[i] [j:j + self.sub_batch_size]) # b, c, h, w fl = model.module.generate_local_fm( images_glb[i:i + 1], patches_var, ratios[i], coordinates[i], [j, j + self.sub_batch_size], len(images), global_model=global_fixed, n_patch_all=len(coordinates[i])) fg = model.module._crop_global( fm_global[i:i + 1], coordinates[i][j:j + self.sub_batch_size], ratios[i])[0] fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear') output_ensembles = model.module.ensemble( fl, fg) # include cordinates # ensemble predictions for pred_ensemble in torch.max( output_ensembles.data, 1)[1].data: predicted_ensembles[i][int( pred_ensemble)] += 1 j += self.sub_batch_size ################################################### _, predictions_global = torch.max(outputs_global.data, 1) if not self.test: self.metrics_global.update(labels_glb, predictions_global) if self.mode == 2 or self.mode == 3: # patch predictions ########################### predictions_local = predicted_patches.argmax(1) if not self.test: self.metrics_local.update(labels_glb, predictions_local) ################################################### predictions = predicted_ensembles.argmax(1) if not self.test: self.metrics.update(labels_glb, predictions) return predictions, predictions_global, predictions_local else: return None, predictions_global, None
class Trainer(object): def __init__(self, criterion, optimizer, n_class, sub_batchsize, mode=1, fmreg=0.15): self.criterion = criterion self.optimizer = optimizer self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class self.sub_batchsize = sub_batchsize self.mode = mode self.fmreg = fmreg # regulization item def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def train(self, sample, model): model.train() labels = sample['label'].squeeze(1).long() labels_npy = np.array(labels) labels_torch = labels.cuda() h, w = sample['output_size'][0] # print(labels[0].size) if self.mode == 1: # global img_g = sample['image_g'].cuda() outputs_g = model.forward(img_g) outputs_g = F.interpolate(outputs_g, size=(h, w), mode='bilinear') # print(outputs_g.size(), labels_torch.size()) loss = self.criterion(outputs_g, labels_torch) loss.backward() self.optimizer.step() self.optimizer.zero_grad() if self.mode == 2: # local img_l = sample['image_l'].cuda() batch_size = img_l.size(0) idx = 0 outputs_l = [] while idx+self.sub_batchsize <= batch_size: output_l = model.forward(img_l[idx:idx+self.sub_batchsize]) output_l = F.interpolate(output_l, size=(h, w), mode='bilinear') outputs_l.append(output_l) loss = self.criterion(output_l, labels_torch[idx:idx+self.sub_batchsize]) loss.backward() idx += self.sub_batchsize outputs_l = torch.cat(outputs_l, dim=0) self.optimizer.step() self.optimizer.zero_grad() if self.mode == 3: # global&local img_g = sample['image_g'].cuda() img_l = sample['image_l'].cuda() batch_size = img_l.size(0) idx = 0 outputs = []; outputs_g = []; outputs_l = [] while idx+self.sub_batchsize <= batch_size: output, output_g, output_l, mse = model.forward(img_g[idx:idx+self.sub_batchsize], img_l[idx:idx+self.sub_batchsize], target=labels_torch[idx:idx+self.sub_batchsize]) outputs.append(output); outputs_g.append(output_g); outputs_l.append(output_l) loss = 2* self.criterion(output, labels_torch[idx:idx+self.sub_batchsize]) + self.criterion(output_g, labels_torch[idx:idx+self.sub_batchsize]) + \ self.criterion(output_l, labels_torch[idx:idx+self.sub_batchsize]) + self.fmreg * mse loss.backward() idx += self.sub_batchsize outputs = torch.cat(outputs, dim=0); outputs_g = torch.cat(outputs_g, dim=0); outputs_l = torch.cat(outputs_l, dim=0) self.optimizer.step() self.optimizer.zero_grad() # predictions if self.mode == 1: outputs_g = outputs_g.cpu() predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] self.metrics_global.update(labels_npy, predictions_global) if self.mode == 2: outputs_l = outputs_l.cpu() predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] self.metrics_local.update(labels_npy, predictions_local) if self.mode == 3: outputs_g = outputs_g.cpu(); outputs_l = outputs_l.cpu(); outputs = outputs.cpu() predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] predictions = [outputs[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] self.metrics_global.update(labels_npy, predictions_global) self.metrics_local.update(labels_npy, predictions_local) self.metrics.update(labels_npy, predictions) return loss
class Evaluator(object): def __init__(self, n_class, sub_batchsize, mode=1, test=False): self.metrics_global = ConfusionMatrix(n_class) self.metrics_local = ConfusionMatrix(n_class) self.metrics = ConfusionMatrix(n_class) self.n_class = n_class self.sub_batchsize = sub_batchsize self.mode = mode self.test = test def get_scores(self): score_train = self.metrics.get_scores() score_train_local = self.metrics_local.get_scores() score_train_global = self.metrics_global.get_scores() return score_train, score_train_global, score_train_local def reset_metrics(self): self.metrics.reset() self.metrics_local.reset() self.metrics_global.reset() def eval_test(self, sample, model): with torch.no_grad(): ids = sample['id'] h, w = sample['output_size'][0] if not self.test: labels = sample['label'].squeeze(1).long() labels_npy = np.array(labels) if self.mode == 1: # global img_g = sample['image_g'].cuda() outputs_g = model.forward(img_g) outputs_g = F.interpolate(outputs_g, size=(h, w), mode='bilinear') if self.mode == 2: # local img_l = sample['image_l'].cuda() batch_size = img_l.size(0) idx = 0 outputs_l = [] while idx+self.sub_batchsize <= batch_size: output_l = model.forward(img_l[idx:idx+self.sub_batchsize]) output_l = F.interpolate(output_l, size=(h, w), mode='bilinear') outputs_l.append(output_l) idx += self.sub_batchsize outputs_l = torch.cat(outputs_l, dim=0) if self.mode == 3: # global&local img_g = sample['image_g'].cuda() img_l = sample['image_l'].cuda() batch_size = img_l.size(0) idx = 0 outputs = []; outputs_g = []; outputs_l = [] while idx+self.sub_batchsize <= batch_size: output, output_g, output_l, mse = model.forward(img_g[idx:idx+self.sub_batchsize], img_l[idx:idx+self.sub_batchsize]) outputs.append(output); outputs_g.append(output_g); outputs_l.append(output_l) idx += self.sub_batchsize outputs = torch.cat(outputs, dim=0); outputs_g = torch.cat(outputs_g, dim=0); outputs_l = torch.cat(outputs_l, dim=0) # no target # outputs, outputs_g, outputs_l, mse = model.forward(img_g, img_l) # predictions if self.mode == 1: outputs_g = outputs_g.cpu() predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] if not self.test: self.metrics_global.update(labels_npy, predictions_global) return None, predictions_global, None if self.mode == 2: outputs_l = outputs_l.cpu() predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] if not self.test: self.metrics_local.update(labels_npy, predictions_local) return None, None, predictions_local if self.mode == 3: outputs_g = outputs_g.cpu(); outputs_l = outputs_l.cpu(); outputs = outputs.cpu() predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] predictions = [outputs[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))] if not self.test: self.metrics_global.update(labels_npy, predictions_global) self.metrics_local.update(labels_npy, predictions_local) self.metrics.update(labels_npy, predictions) return predictions, predictions_global, predictions_local