def create_generator(args, params, input_size): # Visualize anchors to see if they match for per image # For visualization purpose, no argumentation will be applied at here, else you may see distortion. # Except resize op, as you will resize your image when you train it anyway. training_params = {'batch_size': 1, 'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': args.num_workers} val_params = {'batch_size': 1, 'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': args.num_workers} training_set = CocoDataset(root_dir=args.dataset_path, set=params.train_set, transform=transforms.Compose([Resizer(input_size)])) training_generator = DataLoader(training_set, **training_params) val_set = CocoDataset(root_dir=args.dataset_path, set=params.val_set, transform=transforms.Compose([Resizer(input_size)])) val_generator = DataLoader(val_set, **val_params) return training_generator, val_generator
def _get_val_data_loader(args): logger.info("Getting val data loader") dataset = CocoDataset( root_dir=args.data_dir, set="val", transform=transforms.Compose([ Normalizer(mean=args.mean, std=args.std), Resizer(_INPUT_SIZES[args.compound_coef]), ]), ) return torch.utils.data.DataLoader( dataset, batch_size=args.batch_size * num_gpus, shuffle=False, drop_last=True, collate_fn=collater, num_workers=args.num_workers if args.num_workers >= 0 else args.batch_size * num_gpus, )
def train(opt): params = Params(f'projects/{opt.project}.yml') if params.num_gpus == 0: os.environ['CUDA_VISIBLE_DEVICES'] = '-1' if torch.cuda.is_available(): torch.cuda.manual_seed(42) else: torch.manual_seed(42) opt.saved_path = opt.saved_path + f'/{params.project_name}/' opt.log_path = opt.log_path + f'/{params.project_name}/tensorboard/' os.makedirs(opt.log_path, exist_ok=True) os.makedirs(opt.saved_path, exist_ok=True) training_params = { 'batch_size': opt.batch_size, 'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } val_params = { 'batch_size': opt.batch_size, 'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1356] training_set = CocoDataset(root_dir=os.path.join(opt.data_path, params.project_name), set=params.train_set, transform=transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Augmenter(), Resizer(input_sizes[opt.compound_coef]) ])) training_generator = DataLoader(training_set, **training_params) val_set = CocoDataset(root_dir=os.path.join(opt.data_path, params.project_name), set=params.val_set, transform=transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Resizer(input_sizes[opt.compound_coef]) ])) val_generator = DataLoader(val_set, **val_params) model = EfficientDetBackbone(num_classes=len(params.obj_list), compound_coef=opt.compound_coef, ratios=eval(params.anchors_ratios), scales=eval(params.anchors_scales)) # load last weights if opt.load_weights is not None: if opt.load_weights.endswith('.pth'): weights_path = opt.load_weights else: weights_path = get_last_weights(opt.saved_path) try: last_step = int( os.path.basename(weights_path).split('_')[-1].split('.')[0]) except: last_step = 0 try: ret = model.load_state_dict(torch.load(weights_path), strict=False) except RuntimeError as e: print(f'[Warning] Ignoring {e}') print( '[Warning] Don\'t panic if you see this, this might be because you load a pretrained weights with different number of classes. The rest of the weights should be loaded already.' ) print( f'[Info] loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}' ) else: last_step = 0 print('[Info] initializing weights...') init_weights(model) # freeze backbone if train head_only if opt.head_only: def freeze_backbone(m): classname = m.__class__.__name__ for ntl in ['EfficientNet', 'BiFPN']: if ntl in classname: for param in m.parameters(): param.requires_grad = False model.apply(freeze_backbone) print('[Info] freezed backbone') # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch # apply sync_bn when using multiple gpu and batch_size per gpu is lower than 4 # useful when gpu memory is limited. # because when bn is disable, the training will be very unstable or slow to converge, # apply sync_bn can solve it, # by packing all mini-batch across all gpus as one batch and normalize, then send it back to all gpus. # but it would also slow down the training by a little bit. if params.num_gpus > 1 and opt.batch_size // params.num_gpus < 4: model.apply(replace_w_sync_bn) use_sync_bn = True else: use_sync_bn = False writer = SummaryWriter( opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/') # warp the model with loss function, to reduce the memory usage on gpu0 and speedup model = ModelWithLoss(model, debug=opt.debug) if params.num_gpus > 0: model = model.cuda() if params.num_gpus > 1: model = CustomDataParallel(model, params.num_gpus) if use_sync_bn: patch_replication_callback(model) if opt.optim == 'adamw': optimizer = torch.optim.AdamW(model.parameters(), opt.lr) else: optimizer = torch.optim.SGD(model.parameters(), opt.lr, momentum=0.9, nesterov=True) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True) epoch = 0 best_loss = 1e5 best_epoch = 0 step = max(0, last_step) model.train() num_iter_per_epoch = len(training_generator) try: for epoch in range(opt.num_epochs): last_epoch = step // num_iter_per_epoch if epoch < last_epoch: continue epoch_loss = [] progress_bar = tqdm(training_generator) for iter, data in enumerate(progress_bar): if iter < step - last_epoch * num_iter_per_epoch: progress_bar.update() continue try: imgs = data['img'] annot = data['annot'] ################just GT check######################### # img_sample = imgs[0,:,:,:] # annot_sample = annot[0,:,:] # img_out = img_sample.numpy() # img_out = np.transpose(img_out, (1,2,0)) # img_out = cv2.cvtColor(img_out, cv2.COLOR_RGB2BGR) # annot_out = annot_sample.numpy() # count, _ = annot_out.shape # for i in range(count): # if annot_out[i,4] >= 0: # cv2.rectangle(img_out, (int(annot_out[i,0]),int(annot_out[i,1])), (int(annot_out[i,2]),int(annot_out[i,3])), (255,0,0), 1) # # cv2.imwrite("test.png", img_out*255) # ####################################################### if params.num_gpus == 1: # if only one gpu, just send it to cuda:0 # elif multiple gpus, send it to multiple gpus in CustomDataParallel, not here imgs = imgs.cuda() annot = annot.cuda() optimizer.zero_grad() cls_loss, reg_loss = model(imgs, annot, obj_list=params.obj_list) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) optimizer.step() epoch_loss.append(float(loss)) progress_bar.set_description( 'Step: {}. Epoch: {}/{}. Iteration: {}/{}. Cls loss: {:.5f}. Reg loss: {:.5f}. Total loss: {:.5f}' .format(step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch, cls_loss.item(), reg_loss.item(), loss.item())) writer.add_scalars('Loss', {'train': loss}, step) writer.add_scalars('Regression_loss', {'train': reg_loss}, step) writer.add_scalars('Classfication_loss', {'train': cls_loss}, step) # log learning_rate current_lr = optimizer.param_groups[0]['lr'] writer.add_scalar('learning_rate', current_lr, step) step += 1 if step % opt.save_interval == 0 and step > 0: save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth' ) print('checkpoint...') except Exception as e: print('[Error]', traceback.format_exc()) print(e) continue scheduler.step(np.mean(epoch_loss)) if epoch % opt.val_interval == 0: model.eval() loss_regression_ls = [] loss_classification_ls = [] for iter, data in enumerate(val_generator): with torch.no_grad(): imgs = data['img'] annot = data['annot'] if params.num_gpus == 1: imgs = imgs.cuda() annot = annot.cuda() cls_loss, reg_loss = model(imgs, annot, obj_list=params.obj_list) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss_classification_ls.append(cls_loss.item()) loss_regression_ls.append(reg_loss.item()) cls_loss = np.mean(loss_classification_ls) reg_loss = np.mean(loss_regression_ls) loss = cls_loss + reg_loss print( 'Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}' .format(epoch, opt.num_epochs, cls_loss, reg_loss, loss)) writer.add_scalars('Loss', {'val': loss}, step) writer.add_scalars('Regression_loss', {'val': reg_loss}, step) writer.add_scalars('Classfication_loss', {'val': cls_loss}, step) if loss + opt.es_min_delta < best_loss: best_loss = loss best_epoch = epoch save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth' ) model.train() # Early stopping if epoch - best_epoch > opt.es_patience > 0: print( '[Info] Stop training at epoch {}. The lowest loss achieved is {}' .format(epoch, best_loss)) break except KeyboardInterrupt: save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth') writer.close() writer.close()
def train(opt): params = Params(f'projects/{opt.project}.yml') if params.num_gpus == 0: os.environ['CUDA_VISIBLE_DEVICES'] = '-1' if torch.cuda.is_available(): torch.cuda.manual_seed(42) else: torch.manual_seed(42) opt.saved_path = opt.saved_path + f'/{params.project_name}/' opt.log_path = opt.log_path + f'/{params.project_name}/tensorboard/' os.makedirs(opt.log_path, exist_ok=True) os.makedirs(opt.saved_path, exist_ok=True) training_params = { 'batch_size': opt.batch_size, 'shuffle': True, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } val_params = { 'batch_size': opt.batch_size, 'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536] training_set = CocoDataset(root_dir=opt.data_path + params.project_name, set=params.train_set, transform=transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Augmenter(), Resizer(input_sizes[opt.compound_coef]) ])) training_generator = DataLoader(training_set, **training_params) val_set = CocoDataset(root_dir=opt.data_path + params.project_name, set=params.val_set, transform=transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Resizer(input_sizes[opt.compound_coef]) ])) val_generator = DataLoader(val_set, **val_params) model = EfficientDetBackbone(num_anchors=9, num_classes=len(params.obj_list), compound_coef=opt.compound_coef) # load last weights if opt.load_weights is not None: if opt.load_weights.endswith('.pth'): weights_path = opt.load_weights else: weights_path = get_last_weights(opt.saved_path) try: last_step = int( os.path.basename(weights_path).split('_')[-1].split('.')[0]) except: last_step = 0 model.load_state_dict(torch.load(weights_path)) print( f'loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}' ) else: last_step = 0 print('initializing weights...') init_weights(model) # freeze backbone if train head_only if opt.head_only: def freeze_backbone(m): classname = m.__class__.__name__ for ntl in ['EfficientNet', 'BiFPN']: if ntl in classname: for param in m.parameters(): param.requires_grad = False model.apply(freeze_backbone) print('freezed backbone') # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch # apply sync_bn when using multiple gpu and batch_size per gpu is lower than 4 # useful when gpu memory is limited. # because when bn is disable, the training will be very unstable or slow to converge, # apply sync_bn can solve it, # by packing all mini-batch across all gpus as one batch and normalize, then send it back to all gpus. # but it would also slow down the training by a little bit. if params.num_gpus > 1 and opt.batch_size // params.num_gpus < 4: model.apply(replace_w_sync_bn) writer = SummaryWriter( opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/') if params.num_gpus > 0: model = model.cuda() model = CustomDataParallel(model, params.num_gpus) optimizer = torch.optim.AdamW(model.parameters(), opt.lr) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True) criterion = FocalLoss() best_loss = 1e5 best_epoch = 0 step = max(0, last_step) model.train() num_iter_per_epoch = len(training_generator) for epoch in range(opt.num_epochs): try: model.train() epoch_loss = [] progress_bar = tqdm(training_generator) for iter, data in enumerate(progress_bar): try: imgs = data['img'] annot = data['annot'] if params.num_gpus > 0: annot = annot.cuda() optimizer.zero_grad() _, regression, classification, anchors = model(imgs) cls_loss, reg_loss = criterion( classification, regression, anchors, annot, # imgs=imgs, obj_list=params.obj_list # uncomment this to debug ) loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) optimizer.step() epoch_loss.append(float(loss)) progress_bar.set_description( 'Step: {}. Epoch: {}/{}. Iteration: {}/{}. Cls loss: {:.5f}. Reg loss: {:.5f}. Total loss: {:.5f}' .format(step, epoch + 1, opt.num_epochs, iter + 1, num_iter_per_epoch, cls_loss.item(), reg_loss.item(), loss.item())) writer.add_scalars('Loss', {'train': loss}, step) writer.add_scalars('Regression_loss', {'train': reg_loss}, step) writer.add_scalars('Classfication_loss', {'train': cls_loss}, step) # log learning_rate current_lr = optimizer.param_groups[0]['lr'] writer.add_scalar('learning_rate', current_lr, step) step += 1 except Exception as e: print(traceback.format_exc()) print(e) continue scheduler.step(np.mean(epoch_loss)) if step % opt.save_interval == 0 and step > 0: save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth') if epoch % opt.val_interval == 0: model.eval() loss_regression_ls = [] loss_classification_ls = [] for iter, data in enumerate(val_generator): with torch.no_grad(): imgs = data['img'] annot = data['annot'] if params.num_gpus > 0: annot = annot.cuda() _, regression, classification, anchors = model(imgs) cls_loss, reg_loss = criterion(classification, regression, anchors, annot) loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss_classification_ls.append(cls_loss.item()) loss_regression_ls.append(reg_loss.item()) cls_loss = np.mean(loss_classification_ls) reg_loss = np.mean(loss_regression_ls) loss = cls_loss + reg_loss print( 'Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}' .format(epoch + 1, opt.num_epochs, cls_loss, reg_loss, loss.mean())) writer.add_scalars('Total_loss', {'val': loss}, step) writer.add_scalars('Regression_loss', {'val': reg_loss}, step) writer.add_scalars('Classfication_loss', {'val': cls_loss}, step) if loss + opt.es_min_delta < best_loss: best_loss = loss best_epoch = epoch save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth' ) # onnx export is not tested. # dummy_input = torch.rand(opt.batch_size, 3, 512, 512) # if torch.cuda.is_available(): # dummy_input = dummy_input.cuda() # if isinstance(model, nn.DataParallel): # model.module.backbone_net.model.set_swish(memory_efficient=False) # # torch.onnx.export(model.module, dummy_input, # os.path.join(opt.saved_path, 'signatrix_efficientdet_coco.onnx'), # verbose=False) # model.module.backbone_net.model.set_swish(memory_efficient=True) # else: # model.backbone_net.model.set_swish(memory_efficient=False) # # torch.onnx.export(model, dummy_input, # os.path.join(opt.saved_path, 'signatrix_efficientdet_coco.onnx'), # verbose=False) # model.backbone_net.model.set_swish(memory_efficient=True) # Early stopping if epoch - best_epoch > opt.es_patience > 0: print( 'Stop training at epoch {}. The lowest loss achieved is {}' .format(epoch, loss)) break writer.close() except KeyboardInterrupt: save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth')
def start_training(self): if self.system_dict["params"]["num_gpus"] == 0: os.environ['CUDA_VISIBLE_DEVICES'] = '-1' if torch.cuda.is_available(): torch.cuda.manual_seed(42) else: torch.manual_seed(42) self.system_dict["params"]["saved_path"] = self.system_dict["params"][ "saved_path"] + "/" + self.system_dict["params"][ "project_name"] + "/" self.system_dict["params"]["log_path"] = self.system_dict["params"][ "log_path"] + "/" + self.system_dict["params"][ "project_name"] + "/tensorboard/" os.makedirs(self.system_dict["params"]["saved_path"], exist_ok=True) os.makedirs(self.system_dict["params"]["log_path"], exist_ok=True) training_params = { 'batch_size': self.system_dict["params"]["batch_size"], 'shuffle': True, 'drop_last': True, 'collate_fn': collater, 'num_workers': self.system_dict["params"]["num_workers"] } val_params = { 'batch_size': self.system_dict["params"]["batch_size"], 'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': self.system_dict["params"]["num_workers"] } input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536] training_set = CocoDataset( self.system_dict["dataset"]["train"]["root_dir"], self.system_dict["dataset"]["train"]["coco_dir"], self.system_dict["dataset"]["train"]["img_dir"], set_dir=self.system_dict["dataset"]["train"]["set_dir"], transform=transforms.Compose([ Normalizer(mean=self.system_dict["params"]["mean"], std=self.system_dict["params"]["std"]), Augmenter(), Resizer( input_sizes[self.system_dict["params"]["compound_coef"]]) ])) training_generator = DataLoader(training_set, **training_params) if (self.system_dict["dataset"]["val"]["status"]): val_set = CocoDataset( self.system_dict["dataset"]["val"]["root_dir"], self.system_dict["dataset"]["val"]["coco_dir"], self.system_dict["dataset"]["val"]["img_dir"], set_dir=self.system_dict["dataset"]["val"]["set_dir"], transform=transforms.Compose([ Normalizer(self.system_dict["params"]["mean"], self.system_dict["params"]["std"]), Resizer(input_sizes[self.system_dict["params"] ["compound_coef"]]) ])) val_generator = DataLoader(val_set, **val_params) print("") print("") model = EfficientDetBackbone( num_classes=len(self.system_dict["params"]["obj_list"]), compound_coef=self.system_dict["params"]["compound_coef"], ratios=eval(self.system_dict["params"]["anchors_ratios"]), scales=eval(self.system_dict["params"]["anchors_scales"])) os.makedirs("pretrained_weights", exist_ok=True) if (self.system_dict["params"]["compound_coef"] == 0): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d0.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 1): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d1.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 2): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d2.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 3): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d3.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 4): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d4.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 5): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d5.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 6): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d6.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) elif (self.system_dict["params"]["compound_coef"] == 7): if (not os.path.isfile( self.system_dict["params"]["load_weights"])): print("Downloading weights") cmd = "wget https://github.com/zylo117/Yet-Another-Efficient-Pytorch/releases/download/1.0/efficientdet-d7.pth -O " + \ self.system_dict["params"]["load_weights"] os.system(cmd) # load last weights if self.system_dict["params"]["load_weights"] is not None: if self.system_dict["params"]["load_weights"].endswith('.pth'): weights_path = self.system_dict["params"]["load_weights"] else: weights_path = get_last_weights( self.system_dict["params"]["saved_path"]) try: last_step = int( os.path.basename(weights_path).split('_')[-1].split('.') [0]) except: last_step = 0 try: ret = model.load_state_dict(torch.load(weights_path), strict=False) except RuntimeError as e: print(f'[Warning] Ignoring {e}') print( '[Warning] Don\'t panic if you see this, this might be because you load a pretrained weights with different number of classes. The rest of the weights should be loaded already.' ) print( f'[Info] loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}' ) else: last_step = 0 print('[Info] initializing weights...') init_weights(model) print("") print("") # freeze backbone if train head_only if self.system_dict["params"]["head_only"]: def freeze_backbone(m): classname = m.__class__.__name__ for ntl in ['EfficientNet', 'BiFPN']: if ntl in classname: for param in m.parameters(): param.requires_grad = False model.apply(freeze_backbone) print('[Info] freezed backbone') print("") print("") if self.system_dict["params"]["num_gpus"] > 1 and self.system_dict[ "params"]["batch_size"] // self.system_dict["params"][ "num_gpus"] < 4: model.apply(replace_w_sync_bn) use_sync_bn = True else: use_sync_bn = False writer = SummaryWriter( self.system_dict["params"]["log_path"] + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/') model = ModelWithLoss(model, debug=self.system_dict["params"]["debug"]) if self.system_dict["params"]["num_gpus"] > 0: model = model.cuda() if self.system_dict["params"]["num_gpus"] > 1: model = CustomDataParallel( model, self.system_dict["params"]["num_gpus"]) if use_sync_bn: patch_replication_callback(model) if self.system_dict["params"]["optim"] == 'adamw': optimizer = torch.optim.AdamW(model.parameters(), self.system_dict["params"]["lr"]) else: optimizer = torch.optim.SGD(model.parameters(), self.system_dict["params"]["lr"], momentum=0.9, nesterov=True) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True) epoch = 0 best_loss = 1e5 best_epoch = 0 step = max(0, last_step) model.train() num_iter_per_epoch = len(training_generator) try: for epoch in range(self.system_dict["params"]["num_epochs"]): last_epoch = step // num_iter_per_epoch if epoch < last_epoch: continue epoch_loss = [] progress_bar = tqdm(training_generator) for iter, data in enumerate(progress_bar): if iter < step - last_epoch * num_iter_per_epoch: progress_bar.update() continue try: imgs = data['img'] annot = data['annot'] if self.system_dict["params"]["num_gpus"] == 1: # if only one gpu, just send it to cuda:0 # elif multiple gpus, send it to multiple gpus in CustomDataParallel, not here imgs = imgs.cuda() annot = annot.cuda() optimizer.zero_grad() cls_loss, reg_loss = model( imgs, annot, obj_list=self.system_dict["params"]["obj_list"]) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) optimizer.step() epoch_loss.append(float(loss)) progress_bar.set_description( 'Step: {}. Epoch: {}/{}. Iteration: {}/{}. Cls loss: {:.5f}. Reg loss: {:.5f}. Total loss: {:.5f}' .format(step, epoch, self.system_dict["params"]["num_epochs"], iter + 1, num_iter_per_epoch, cls_loss.item(), reg_loss.item(), loss.item())) writer.add_scalars('Loss', {'train': loss}, step) writer.add_scalars('Regression_loss', {'train': reg_loss}, step) writer.add_scalars('Classfication_loss', {'train': cls_loss}, step) # log learning_rate current_lr = optimizer.param_groups[0]['lr'] writer.add_scalar('learning_rate', current_lr, step) step += 1 if step % self.system_dict["params"][ "save_interval"] == 0 and step > 0: self.save_checkpoint( model, f'efficientdet-d{self.system_dict["params"]["compound_coef"]}_trained.pth' ) #print('checkpoint...') except Exception as e: print('[Error]', traceback.format_exc()) print(e) continue scheduler.step(np.mean(epoch_loss)) if epoch % self.system_dict["params"][ "val_interval"] == 0 and self.system_dict["dataset"][ "val"]["status"]: print("Running validation") model.eval() loss_regression_ls = [] loss_classification_ls = [] for iter, data in enumerate(val_generator): with torch.no_grad(): imgs = data['img'] annot = data['annot'] if self.system_dict["params"]["num_gpus"] == 1: imgs = imgs.cuda() annot = annot.cuda() cls_loss, reg_loss = model( imgs, annot, obj_list=self.system_dict["params"] ["obj_list"]) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss_classification_ls.append(cls_loss.item()) loss_regression_ls.append(reg_loss.item()) cls_loss = np.mean(loss_classification_ls) reg_loss = np.mean(loss_regression_ls) loss = cls_loss + reg_loss print( 'Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}' .format(epoch, self.system_dict["params"]["num_epochs"], cls_loss, reg_loss, loss)) writer.add_scalars('Loss', {'val': loss}, step) writer.add_scalars('Regression_loss', {'val': reg_loss}, step) writer.add_scalars('Classfication_loss', {'val': cls_loss}, step) if loss + self.system_dict["params"][ "es_min_delta"] < best_loss: best_loss = loss best_epoch = epoch self.save_checkpoint( model, f'efficientdet-d{self.system_dict["params"]["compound_coef"]}_trained.pth' ) model.train() # Early stopping if epoch - best_epoch > self.system_dict["params"][ "es_patience"] > 0: print( '[Info] Stop training at epoch {}. The lowest loss achieved is {}' .format(epoch, best_loss)) break except KeyboardInterrupt: self.save_checkpoint( model, f'efficientdet-d{self.system_dict["params"]["compound_coef"]}_trained.pth' ) writer.close() writer.close() print("") print("") print("Training complete")
'shuffle': False, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } # tf bilinear interpolation is different from any other's, just make do input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536] input_size = input_sizes[ opt. compound_coef] if opt.force_input_size is None else opt.force_input_size val_set = CocoDataset(root_dir=os.path.join(opt.data_path, params.project_name), set=params.val_set, transform=transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Resizer(input_sizes[opt.compound_coef]) ])) val_generator = DataLoader(val_set, **val_params) model = EfficientDetBackbone(compound_coef=opt.compound_coef, num_classes=len(obj_list), ratios=anchor_ratios, scales=anchor_scales) model.load_state_dict( torch.load( f'logs/{opt.project}/efficientdet-d{opt.compound_coef}_{opt.weights}.pth', map_location='cpu')) model.requires_grad_(False) model.eval()
def train(opt): ''' Input: get_args() Function: Train the model. ''' params = Params(f'projects/{opt.project}.yml') if params.num_gpus == 0: os.environ['CUDA_VISIBLE_DEVICES'] = '-1' if torch.cuda.is_available(): torch.cuda.manual_seed(42) else: torch.manual_seed(42) opt.saved_path = opt.saved_path + f'/{params.project_name}/' opt.log_path = opt.log_path + f'/{params.project_name}/tensorboard/' os.makedirs(opt.log_path, exist_ok=True) os.makedirs(opt.saved_path, exist_ok=True) # evaluation json file pred_folder = f'{OPT.data_path}/{OPT.project}/predictions' os.makedirs(pred_folder, exist_ok=True) evaluation_pred_file = f'{pred_folder}/instances_bbox_results.json' training_params = { 'batch_size': opt.batch_size, 'shuffle': True, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } val_params = { 'batch_size': opt.batch_size, 'shuffle': True, 'drop_last': True, 'collate_fn': collater, 'num_workers': opt.num_workers } input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536] training_set = CocoDataset(root_dir=os.path.join(opt.data_path, params.project_name), set=params.train_set, transform=torchvision.transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Augmenter(), Resizer(input_sizes[opt.compound_coef]) ])) training_generator = DataLoader(training_set, **training_params) val_set = CocoDataset(root_dir=os.path.join(opt.data_path, params.project_name), set=params.val_set, transform=torchvision.transforms.Compose([ Normalizer(mean=params.mean, std=params.std), Resizer(input_sizes[opt.compound_coef]) ])) val_generator = DataLoader(val_set, **val_params) model = EfficientDetBackbone(num_classes=len(params.obj_list), compound_coef=opt.compound_coef, ratios=eval(params.anchors_ratios), scales=eval(params.anchors_scales)) # load last weights if opt.load_weights is not None: if opt.load_weights.endswith('.pth'): weights_path = opt.load_weights else: weights_path = get_last_weights(opt.saved_path) try: last_step = int( os.path.basename(weights_path).split('_')[-1].split('.')[0]) except Exception as exception: last_step = 0 try: _ = model.load_state_dict(torch.load(weights_path), strict=False) except RuntimeError as rerror: print(f'[Warning] Ignoring {rerror}') print('[Warning] Don\'t panic if you see this, '\ 'this might be because you load a pretrained weights with different number of classes.'\ ' The rest of the weights should be loaded already.') print( f'[Info] loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}' ) else: last_step = 0 print('[Info] initializing weights...') init_weights(model) # freeze backbone if train head_only if opt.head_only: def freeze_backbone(mdl): classname = mdl.__class__.__name__ for ntl in ['EfficientNet', 'BiFPN']: if ntl in classname: for param in mdl.parameters(): param.requires_grad = False model.apply(freeze_backbone) print('[Info] freezed backbone') # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch # apply sync_bn when using multiple gpu and batch_size per gpu is lower than 4 # useful when gpu memory is limited. # because when bn is disable, the training will be very unstable or slow to converge, # apply sync_bn can solve it, # by packing all mini-batch across all gpus as one batch and normalize, then send it back to all gpus. # but it would also slow down the training by a little bit. if params.num_gpus > 1 and opt.batch_size // params.num_gpus < 4: model.apply(replace_w_sync_bn) use_sync_bn = True else: use_sync_bn = False writer = SummaryWriter( opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/') # warp the model with loss function, to reduce the memory usage on gpu0 and speedup model = ModelWithLoss(model, debug=opt.debug) if params.num_gpus > 0: model = model.cuda() if params.num_gpus > 1: model = CustomDataParallel(model, params.num_gpus) if use_sync_bn: patch_replication_callback(model) if opt.optim == 'adamw': optimizer = torch.optim.AdamW(model.parameters(), opt.lr) else: optimizer = torch.optim.SGD(model.parameters(), opt.lr, momentum=0.9, nesterov=True) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True) epoch = 0 best_loss = 1e5 best_epoch = 0 step = max(0, last_step) model.train() num_iter_per_epoch = len(training_generator) num_val_iter_per_epoch = len(val_generator) # Limit the no.of preds to #images in val. # Here, I averaged the #obj to 5 for computational efficacy if opt.max_preds_toeval > 0: opt.max_preds_toeval = len(val_generator) * opt.batch_size * 5 try: for epoch in range(opt.num_epochs): last_epoch = step // num_iter_per_epoch if epoch < last_epoch: continue epoch_loss = [] progress_bar = tqdm(training_generator) for iternum, data in enumerate(progress_bar): if iternum < step - last_epoch * num_iter_per_epoch: progress_bar.update() continue try: imgs = data['img'] annot = data['annot'] if params.num_gpus == 1: # if only one gpu, just send it to cuda:0 # elif multiple gpus, send it to multiple gpus in CustomDataParallel, not here imgs = imgs.cuda() annot = annot.cuda() optimizer.zero_grad() if iternum % int(num_iter_per_epoch * (opt.eval_percent_epoch / 100)) != 0: model.debug = False cls_loss, reg_loss, _ = model(imgs, annot, obj_list=params.obj_list) else: model.debug = True cls_loss, reg_loss, imgs_labelled = model( imgs, annot, obj_list=params.obj_list) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) optimizer.step() epoch_loss.append(float(loss)) progress_bar.set_description( 'Step: {}. Epoch: {}/{}. Iteration: {}/{}. Cls loss: {:.5f}. Reg loss: {:.5f}. Total loss: {:.5f}' .format(step, epoch, opt.num_epochs, iternum + 1, num_iter_per_epoch, cls_loss.item(), reg_loss.item(), loss.item())) writer.add_scalars('Loss', {'train': loss}, step) writer.add_scalars('Regression_loss', {'train': reg_loss}, step) writer.add_scalars('Classfication_loss', {'train': cls_loss}, step) if iternum % int( num_iter_per_epoch * (opt.eval_percent_epoch / 100)) == 0 and step > 0: # create grid of images imgs_labelled = np.asarray(imgs_labelled) imgs_labelled = torch.from_numpy( imgs_labelled) # (N, H, W, C) imgs_labelled.transpose_(1, 3) # (N, C, H, W) imgs_labelled.transpose_(2, 3) img_grid = torchvision.utils.make_grid(imgs_labelled) # write to tensorboard writer.add_image('Training_images', img_grid, global_step=step) #########################################################start EVAL##################################################### model.eval() model.debug = False # Don't print images in tensorboard now. # remove json if os.path.exists(evaluation_pred_file): os.remove(evaluation_pred_file) loss_regression_ls = [] loss_classification_ls = [] model.evalresults = [ ] # Empty the results for next evaluation. imgs_to_viz = [] num_validation_steps = int( num_val_iter_per_epoch * (opt.eval_sampling_percent / 100)) for valiternum, valdata in enumerate(val_generator): with torch.no_grad(): imgs = valdata['img'] annot = valdata['annot'] resizing_imgs_scales = valdata['scale'] new_ws = valdata['new_w'] new_hs = valdata['new_h'] imgs_ids = valdata['img_id'] if params.num_gpus >= 1: imgs = imgs.cuda() annot = annot.cuda() if valiternum % (num_validation_steps // (opt.num_visualize_images // opt.batch_size)) != 0: model.debug = False cls_loss, reg_loss, _ = model( imgs, annot, obj_list=params.obj_list, resizing_imgs_scales= resizing_imgs_scales, new_ws=new_ws, new_hs=new_hs, imgs_ids=imgs_ids) else: model.debug = True cls_loss, reg_loss, val_imgs_labelled = model( imgs, annot, obj_list=params.obj_list, resizing_imgs_scales= resizing_imgs_scales, new_ws=new_ws, new_hs=new_hs, imgs_ids=imgs_ids) imgs_to_viz += list(val_imgs_labelled) loss_classification_ls.append(cls_loss.item()) loss_regression_ls.append(reg_loss.item()) if valiternum > (num_validation_steps): break cls_loss = np.mean(loss_classification_ls) reg_loss = np.mean(loss_regression_ls) loss = cls_loss + reg_loss print( 'Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}' .format(epoch, opt.num_epochs, cls_loss, reg_loss, loss)) writer.add_scalars('Loss', {'val': loss}, step) writer.add_scalars('Regression_loss', {'val': reg_loss}, step) writer.add_scalars('Classfication_loss', {'val': cls_loss}, step) # create grid of images val_imgs_labelled = np.asarray(imgs_to_viz) val_imgs_labelled = torch.from_numpy( val_imgs_labelled) # (N, H, W, C) val_imgs_labelled.transpose_(1, 3) # (N, C, H, W) val_imgs_labelled.transpose_(2, 3) val_img_grid = torchvision.utils.make_grid( val_imgs_labelled, nrow=2) # write to tensorboard writer.add_image('Eval_Images', val_img_grid, \ global_step=(step)) if opt.max_preds_toeval > 0: json.dump(model.evalresults, open(evaluation_pred_file, 'w'), indent=4) try: val_results = calc_mAP_fin(params.project_name,\ params.val_set, evaluation_pred_file, \ val_gt=f'{OPT.data_path}/{OPT.project}/annotations/instances_{params.val_set}.json') for catgname in val_results: metricname = 'Average Precision (AP) @[ IoU = 0.50 | area = all | maxDets = 100 ]' evalscore = val_results[catgname][ metricname] writer.add_scalars( f'mAP@IoU=0.5 and area=all', {f'{catgname}': evalscore}, step) except Exception as exption: print("Unable to perform evaluation", exption) if loss + opt.es_min_delta < best_loss: best_loss = loss best_epoch = epoch save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth' ) model.train() # Early stopping if epoch - best_epoch > opt.es_patience > 0: print( '[Info] Stop training at epoch {}. The lowest loss achieved is {}' .format(epoch, best_loss)) break #########################################################EVAL##################################################### # log learning_rate current_lr = optimizer.param_groups[0]['lr'] writer.add_scalar('learning_rate', current_lr, step) step += 1 if step % opt.save_interval == 0 and step > 0: save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth' ) print('checkpoint...') except Exception as exception: print('[Error]', traceback.format_exc()) print(exception) continue scheduler.step(np.mean(epoch_loss)) except KeyboardInterrupt: save_checkpoint( model, f'efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth') writer.close() writer.close()
def train(opt): params = Params(opt.config) if params.num_gpus == 0: os.environ["CUDA_VISIBLE_DEVICES"] = "-1" if torch.cuda.is_available(): torch.cuda.manual_seed(42) else: torch.manual_seed(42) opt.saved_path = params.logdir opt.log_path = os.path.join(params.logdir, "tensorboard") os.makedirs(opt.saved_path, exist_ok=True) os.makedirs(opt.log_path, exist_ok=True) training_params = { "batch_size": opt.batch_size, "shuffle": True, "drop_last": True, "collate_fn": collater, "num_workers": opt.num_workers, } val_params = { "batch_size": opt.batch_size, "shuffle": False, "drop_last": True, "collate_fn": collater, "num_workers": opt.num_workers, } input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536] training_set = CocoDataset( image_dir=params.image_dir, json_path=params.train_annotations, transform=transforms.Compose( [ Normalizer(mean=params.mean, std=params.std), Augmenter(), Resizer(input_sizes[opt.compound_coef]), ] ), ) training_generator = DataLoader(training_set, **training_params) if params.val_image_dir is None: params.val_image_dir = params.image_dir val_set = CocoDataset( image_dir=params.val_image_dir, json_path=params.val_annotations, transform=transforms.Compose( [Normalizer(mean=params.mean, std=params.std), Resizer(input_sizes[opt.compound_coef])] ), ) val_generator = DataLoader(val_set, **val_params) model = EfficientDetBackbone( num_classes=len(params.obj_list), compound_coef=opt.compound_coef, ratios=eval(params.anchors_ratios), scales=eval(params.anchors_scales), ) # load last weights if opt.load_weights is not None: if opt.load_weights.endswith(".pth"): weights_path = opt.load_weights else: weights_path = get_last_weights(opt.saved_path) try: last_step = int(os.path.basename(weights_path).split("_")[-1].split(".")[0]) except: last_step = 0 try: ret = model.load_state_dict(torch.load(weights_path), strict=False) except RuntimeError as e: print(f"[Warning] Ignoring {e}") print( "[Warning] Don't panic if you see this, this might be because you load a pretrained weights with different number of classes. The rest of the weights should be loaded already." ) print( f"[Info] loaded weights: {os.path.basename(weights_path)}, resuming checkpoint from step: {last_step}" ) else: last_step = 0 print("[Info] initializing weights...") init_weights(model) # freeze backbone if train head_only if opt.head_only: def freeze_backbone(m): classname = m.__class__.__name__ for ntl in ["EfficientNet", "BiFPN"]: if ntl in classname: for param in m.parameters(): param.requires_grad = False model.apply(freeze_backbone) print("[Info] freezed backbone") # https://github.com/vacancy/Synchronized-BatchNorm-PyTorch # apply sync_bn when using multiple gpu and batch_size per gpu is lower than 4 # useful when gpu memory is limited. # because when bn is disable, the training will be very unstable or slow to converge, # apply sync_bn can solve it, # by packing all mini-batch across all gpus as one batch and normalize, then send it back to all gpus. # but it would also slow down the training by a little bit. if params.num_gpus > 1 and opt.batch_size // params.num_gpus < 4: model.apply(replace_w_sync_bn) use_sync_bn = True else: use_sync_bn = False writer = SummaryWriter(opt.log_path + f'/{datetime.datetime.now().strftime("%Y%m%d-%H%M%S")}/') # warp the model with loss function, to reduce the memory usage on gpu0 and speedup model = ModelWithLoss(model, debug=opt.debug) if params.num_gpus > 0: model = model.cuda() if params.num_gpus > 1: model = CustomDataParallel(model, params.num_gpus) if use_sync_bn: patch_replication_callback(model) if opt.optim == "adamw": optimizer = torch.optim.AdamW(model.parameters(), opt.lr) else: optimizer = torch.optim.SGD(model.parameters(), opt.lr, momentum=0.9, nesterov=True) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, verbose=True) epoch = 0 best_loss = 1e5 best_epoch = 0 step = max(0, last_step) model.train() num_iter_per_epoch = len(training_generator) try: for epoch in range(opt.num_epochs): last_epoch = step // num_iter_per_epoch if epoch < last_epoch: continue epoch_loss = [] progress_bar = tqdm(training_generator) for iter, data in enumerate(progress_bar): if iter < step - last_epoch * num_iter_per_epoch: progress_bar.update() continue try: imgs = data["img"] annot = data["annot"] if params.num_gpus == 1: # if only one gpu, just send it to cuda:0 # elif multiple gpus, send it to multiple gpus in CustomDataParallel, not here imgs = imgs.cuda() annot = annot.cuda() optimizer.zero_grad() cls_loss, reg_loss = model(imgs, annot, obj_list=params.obj_list) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss.backward() # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) optimizer.step() epoch_loss.append(float(loss)) progress_bar.set_description( "Step: {}. Epoch: {}/{}. Iteration: {}/{}. Cls loss: {:.5f}. Reg loss: {:.5f}. Total loss: {:.5f}".format( step, epoch, opt.num_epochs, iter + 1, num_iter_per_epoch, cls_loss.item(), reg_loss.item(), loss.item(), ) ) writer.add_scalars("Loss", {"train": loss}, step) writer.add_scalars("Regression_loss", {"train": reg_loss}, step) writer.add_scalars("Classfication_loss", {"train": cls_loss}, step) # log learning_rate current_lr = optimizer.param_groups[0]["lr"] writer.add_scalar("learning_rate", current_lr, step) step += 1 if step % opt.save_interval == 0 and step > 0: save_checkpoint( model, f"efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth" ) print("checkpoint...") except Exception as e: print("[Error]", traceback.format_exc()) print(e) continue scheduler.step(np.mean(epoch_loss)) if epoch % opt.val_interval == 0: model.eval() loss_regression_ls = [] loss_classification_ls = [] for iter, data in enumerate(val_generator): with torch.no_grad(): imgs = data["img"] annot = data["annot"] if params.num_gpus == 1: imgs = imgs.cuda() annot = annot.cuda() cls_loss, reg_loss = model(imgs, annot, obj_list=params.obj_list) cls_loss = cls_loss.mean() reg_loss = reg_loss.mean() loss = cls_loss + reg_loss if loss == 0 or not torch.isfinite(loss): continue loss_classification_ls.append(cls_loss.item()) loss_regression_ls.append(reg_loss.item()) cls_loss = np.mean(loss_classification_ls) reg_loss = np.mean(loss_regression_ls) loss = cls_loss + reg_loss print( "Val. Epoch: {}/{}. Classification loss: {:1.5f}. Regression loss: {:1.5f}. Total loss: {:1.5f}".format( epoch, opt.num_epochs, cls_loss, reg_loss, loss ) ) writer.add_scalars("Loss", {"val": loss}, step) writer.add_scalars("Regression_loss", {"val": reg_loss}, step) writer.add_scalars("Classfication_loss", {"val": cls_loss}, step) if loss + opt.es_min_delta < best_loss: best_loss = loss best_epoch = epoch save_checkpoint(model, f"efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth") model.train() # Early stopping if epoch - best_epoch > opt.es_patience > 0: print( "[Info] Stop training at epoch {}. The lowest loss achieved is {}".format( epoch, best_loss ) ) break except KeyboardInterrupt: save_checkpoint(model, f"efficientdet-d{opt.compound_coef}_{epoch}_{step}.pth") writer.close() writer.close()
input_sizes = [512, 640, 768, 896, 1024, 1280, 1280, 1536, 1536] compound_coef = 3 project_name = "fujiseal_stain" data_path = "./datasets/" train = True params = Params(f"projects/{project_name}.yml") if train: dataset_name = params.train_set else: dataset_name = params.val_set image_dataset = CocoDataset( root_dir=os.path.join(data_path, params.project_name), set=dataset_name, transform=transforms.Compose([Resizer(input_sizes[compound_coef])])) image_loader = DataLoader(image_dataset, batch_size=1, shuffle=False, drop_last=False, num_workers=0) psum = torch.tensor([0.0, 0.0, 0.0]) psum_sq = torch.tensor([0.0, 0.0, 0.0]) start_time = time.time() for inputs in image_loader: imgs = inputs['img'] psum += imgs.sum(axis=[0, 1, 2])