class NetworkFactory(object): def __init__(self, db): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel( self.network, chunk_sizes=system_configs.chunk_sizes).cuda() self.load_cropped_pretrained_model( "cache/nnet/CenterNet-52/CenterNet-52_480000.pkl") total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) # self.fix_layers() # fix kps and prelayer if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def load_cropped_pretrained_model(self, params_file): x = torch.load(params_file) params = { 'module.model.%s' % k: v for k, v in x.items() if 'heats' not in k } self.network.load_state_dict(params, strict=False) print("load the cropped weights from COCO successfully.") def fix_layers(self): for m, v in self.network.named_parameters(): if '.pre' in m or '.kps' in m: v.requires_grad = False def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [x for x in xs] ys = [y for y in ys] self.optimizer.zero_grad() loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() focal_loss = focal_loss.mean() pull_loss = pull_loss.mean() push_loss = push_loss.mean() regr_loss = regr_loss.mean() loss.backward() self.optimizer.step() return loss, focal_loss, pull_loss, push_loss, regr_loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, db): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.model_name) print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel(self.network, chunk_sizes=system_configs.chunk_sizes) total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [[x.cuda(non_blocking=True) for x in z] for z in xs] ys = [[y.cuda(non_blocking=True) for y in z] for z in ys] self.optimizer.zero_grad() loss = self.network(xs, ys) loss = loss.mean() loss.backward() self.optimizer.step() return loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [[x.cuda(non_blocking=True) for x in z] for z in xs] ys = [[y.cuda(non_blocking=True) for y in z] for z in ys] loss = self.network(xs, ys) loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [[x.cuda(non_blocking=True) for x in z] for z in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, db): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel(self.network, chunk_sizes=system_configs.chunk_sizes) total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] self.optimizer.zero_grad() loss = self.network(xs, ys) loss = loss.mean() loss.backward() self.optimizer.step() return loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss = self.network(xs, ys) loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading pretrain from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_my_state_dict(params, strict=False) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f) def calculate_bboxes(self, cfg_file, iteration): print("JOEJOE") with open(cfg_file, "r") as f: configs = json.load(f) train_split = system_configs.train_split val_split = system_configs.val_split test_split = system_configs.test_split split = { "training": train_split, "validation": val_split, "testing": test_split }["validation"] dataset = system_configs.dataset testing_db = datasets[dataset](configs["db"], split) test_file = "test.{}".format(testing_db.data) testing = importlib.import_module(test_file).testing result_dir = system_configs.result_dir result_dir = os.path.join(result_dir, str(iteration), split) if not os.path.exists(result_dir): os.makedirs(result_dir) testing(testing_db, self, result_dir, iteration)
class NetworkFactory(object): def __init__(self, db): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel( self.network, chunk_sizes=system_configs.chunk_sizes).cuda() total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [x for x in xs] ys = [y for y in ys] self.optimizer.zero_grad() loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() focal_loss = focal_loss.mean() pull_loss = pull_loss.mean() push_loss = push_loss.mean() regr_loss = regr_loss.mean() loss.backward() self.optimizer.step() return loss, focal_loss, pull_loss, push_loss, regr_loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) # modify_params = {} # for k,v in params.items(): # # if k.split('.')[-1] != 'num_batches_tracked': # if (k.split('.')[-1] != 'num_batches_tracked') and (k.split('.')[1] not in ['tl_heats', 'br_heats', 'ct_heats']): # modify_params[k] = v # model_dict = self.model.state_dict() # modify_params = {k: v for k, v in modify_params.items() if k in model_dict} # model_dict.update(modify_params) # # self.model.load_state_dict(modify_params) # self.model.load_state_dict(model_dict) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) # modify_params = {} # for k,v in params.items(): # if (k.split('.')[-1] != 'num_batches_tracked') and (k.split('.')[1] not in ['tl_heats', 'br_heats', 'ct_heats']): # modify_params[k] = v # model_dict = self.model.state_dict() # modify_params = {k: v for k, v in modify_params.items() if k in model_dict} # model_dict.update(modify_params) # # self.model.load_state_dict(modify_params) # self.model.load_state_dict(model_dict) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, db): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel( self.network, chunk_sizes=system_configs.chunk_sizes).cuda() total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [x for x in xs] ys = [y for y in ys] self.optimizer.zero_grad() loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() focal_loss = focal_loss.mean() pull_loss = pull_loss.mean() push_loss = push_loss.mean() regr_loss = regr_loss.mean() loss.backward() self.optimizer.step() return loss, focal_loss, pull_loss, push_loss, regr_loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f) def _rescale_dets(self, detections, ratios, borders, sizes): xs, ys = detections[..., 0:4:2], detections[..., 1:4:2] xs /= ratios[:, 1][:, None, None] ys /= ratios[:, 0][:, None, None] xs -= borders[:, 2][:, None, None] ys -= borders[:, 0][:, None, None] tx_inds = xs[:, :, 0] <= -5 bx_inds = xs[:, :, 1] >= sizes[0, 1] + 5 ty_inds = ys[:, :, 0] <= -5 by_inds = ys[:, :, 1] >= sizes[0, 0] + 5 np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs) np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys) detections[:, tx_inds[0, :], 4] = -1 detections[:, bx_inds[0, :], 4] = -1 detections[:, ty_inds[0, :], 4] = -1 detections[:, by_inds[0, :], 4] = -1 def kp_decode(self, images, K, ae_threshold=0.5, kernel=3): detections, center = self.test([images], ae_threshold=ae_threshold, K=K, kernel=kernel) detections = detections.data.cpu().numpy() center = center.data.cpu().numpy() return detections, center def kp_detection(self, image, db, result_dir, debug=False): K = db.configs["top_k"] ae_threshold = db.configs["ae_threshold"] nms_kernel = db.configs["nms_kernel"] weight_exp = db.configs["weight_exp"] merge_bbox = db.configs["merge_bbox"] categories = db.configs["categories"] nms_threshold = db.configs["nms_threshold"] max_per_image = db.configs["max_per_image"] nms_algorithm = { "nms": 0, "linear_soft_nms": 1, "exp_soft_nms": 2 }[db.configs["nms_algorithm"]] top_bboxes = {} if True: #db_ind = db_inds[ind] image_id = 0 height, width = image.shape[0:2] detections = [] center_points = [] if True: scale = 1 new_height = int(height * scale) new_width = int(width * scale) new_center = np.array([new_height // 2, new_width // 2]) inp_height = new_height | 127 inp_width = new_width | 127 images = np.zeros((1, 3, inp_height, inp_width), dtype=np.float32) ratios = np.zeros((1, 2), dtype=np.float32) borders = np.zeros((1, 4), dtype=np.float32) sizes = np.zeros((1, 2), dtype=np.float32) out_height, out_width = (inp_height + 1) // 4, (inp_width + 1) // 4 height_ratio = out_height / inp_height width_ratio = out_width / inp_width resized_image = cv2.resize(image, (new_width, new_height)) resized_image, border, offset = crop_image( resized_image, new_center, [inp_height, inp_width]) resized_image = resized_image / 255. normalize_(resized_image, db.mean, db.std) images[0] = resized_image.transpose((2, 0, 1)) borders[0] = border sizes[0] = [int(height * scale), int(width * scale)] ratios[0] = [height_ratio, width_ratio] images = np.concatenate((images, images[:, :, :, ::-1]), axis=0) images = torch.from_numpy(images) dets, center = self.kp_decode(images, K, ae_threshold=ae_threshold, kernel=nms_kernel) dets = dets.reshape(2, -1, 8) center = center.reshape(2, -1, 4) dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]] center[1, :, [0]] = out_width - center[1, :, [0]] dets = dets.reshape(1, -1, 8) center = center.reshape(1, -1, 4) self._rescale_dets(dets, ratios, borders, sizes) center[..., [0]] /= ratios[:, 1][:, None, None] center[..., [1]] /= ratios[:, 0][:, None, None] center[..., [0]] -= borders[:, 2][:, None, None] center[..., [1]] -= borders[:, 0][:, None, None] np.clip(center[..., [0]], 0, sizes[:, 1][:, None, None], out=center[..., [0]]) np.clip(center[..., [1]], 0, sizes[:, 0][:, None, None], out=center[..., [1]]) dets[:, :, 0:4] /= scale center[:, :, 0:2] /= scale if scale == 1: center_points.append(center) detections.append(dets) detections = np.concatenate(detections, axis=1) center_points = np.concatenate(center_points, axis=1) classes = detections[..., -1] classes = classes[0] detections = detections[0] center_points = center_points[0] valid_ind = detections[:, 4] > -1 valid_detections = detections[valid_ind] box_width = valid_detections[:, 2] - valid_detections[:, 0] box_height = valid_detections[:, 3] - valid_detections[:, 1] s_ind = (box_width * box_height <= 22500) l_ind = (box_width * box_height > 22500) s_detections = valid_detections[s_ind] l_detections = valid_detections[l_ind] s_left_x = (2 * s_detections[:, 0] + s_detections[:, 2]) / 3 s_right_x = (s_detections[:, 0] + 2 * s_detections[:, 2]) / 3 s_top_y = (2 * s_detections[:, 1] + s_detections[:, 3]) / 3 s_bottom_y = (s_detections[:, 1] + 2 * s_detections[:, 3]) / 3 s_temp_score = copy.copy(s_detections[:, 4]) s_detections[:, 4] = -1 center_x = center_points[:, 0][:, np.newaxis] center_y = center_points[:, 1][:, np.newaxis] s_left_x = s_left_x[np.newaxis, :] s_right_x = s_right_x[np.newaxis, :] s_top_y = s_top_y[np.newaxis, :] s_bottom_y = s_bottom_y[np.newaxis, :] ind_lx = (center_x - s_left_x) > 0 ind_rx = (center_x - s_right_x) < 0 ind_ty = (center_y - s_top_y) > 0 ind_by = (center_y - s_bottom_y) < 0 ind_cls = (center_points[:, 2][:, np.newaxis] - s_detections[:, -1][np.newaxis, :]) == 0 ind_s_new_score = np.max( ((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0)), axis=0) == 1 index_s_new_score = np.argmax( ((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0))[:, ind_s_new_score], axis=0) s_detections[:, 4][ind_s_new_score] = ( s_temp_score[ind_s_new_score] * 2 + center_points[index_s_new_score, 3]) / 3 l_left_x = (3 * l_detections[:, 0] + 2 * l_detections[:, 2]) / 5 l_right_x = (2 * l_detections[:, 0] + 3 * l_detections[:, 2]) / 5 l_top_y = (3 * l_detections[:, 1] + 2 * l_detections[:, 3]) / 5 l_bottom_y = (2 * l_detections[:, 1] + 3 * l_detections[:, 3]) / 5 l_temp_score = copy.copy(l_detections[:, 4]) l_detections[:, 4] = -1 center_x = center_points[:, 0][:, np.newaxis] center_y = center_points[:, 1][:, np.newaxis] l_left_x = l_left_x[np.newaxis, :] l_right_x = l_right_x[np.newaxis, :] l_top_y = l_top_y[np.newaxis, :] l_bottom_y = l_bottom_y[np.newaxis, :] ind_lx = (center_x - l_left_x) > 0 ind_rx = (center_x - l_right_x) < 0 ind_ty = (center_y - l_top_y) > 0 ind_by = (center_y - l_bottom_y) < 0 ind_cls = (center_points[:, 2][:, np.newaxis] - l_detections[:, -1][np.newaxis, :]) == 0 ind_l_new_score = np.max( ((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0)), axis=0) == 1 index_l_new_score = np.argmax( ((ind_lx + 0) & (ind_rx + 0) & (ind_ty + 0) & (ind_by + 0) & (ind_cls + 0))[:, ind_l_new_score], axis=0) l_detections[:, 4][ind_l_new_score] = ( l_temp_score[ind_l_new_score] * 2 + center_points[index_l_new_score, 3]) / 3 detections = np.concatenate([l_detections, s_detections], axis=0) detections = detections[np.argsort(-detections[:, 4])] classes = detections[..., -1] keep_inds = (detections[:, 4] > -1) detections = detections[keep_inds] classes = classes[keep_inds] top_bboxes[image_id] = {} for j in range(categories): keep_inds = (classes == j) top_bboxes[image_id][j + 1] = detections[keep_inds][:, 0:7].astype( np.float32) if merge_bbox: soft_nms_merge(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm, weight_exp=weight_exp) else: soft_nms(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm) top_bboxes[image_id][j + 1] = top_bboxes[image_id][j + 1][:, 0:5] scores = np.hstack([ top_bboxes[image_id][j][:, -1] for j in range(1, categories + 1) ]) if len(scores) > max_per_image: kth = len(scores) - max_per_image thresh = np.partition(scores, kth)[kth] for j in range(1, categories + 1): keep_inds = (top_bboxes[image_id][j][:, -1] >= thresh) top_bboxes[image_id][j] = top_bboxes[image_id][j][ keep_inds] return top_bboxes[image_id] return 0
class NetworkFactory( object): # https://blog.csdn.net/ying86615791/article/details/89531974 def __init__(self, db): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) print("module_file: {}".format( module_file)) # module_file: models.CenterNet-52 nnet_module = importlib.import_module( module_file ) # 导入models.CenterNet-52; NetworkFactory又通过importlib.import_module(module_file)导入了self.model和self.loss # module_file使用的system_configs.snapshot_name来自train.py中的configs["system"]["snapshot_name"] = args.cfg_file # NetworkFactory中的self.model和self.loss, 这二者来自CenterNet-52.py中的class model(kp), 这个model继承自kp.py中的class kp(nn.Module), 这个loss也是来自kp.py中的class AELoss(nn.Module) # 所以model主要框架都在这个class kp(nn.Module)里. # 只传入1张图片的list时, 模型执行_test函数. 所以在测试的时候(看test/coco.py中的def kp_decode函数), 输入被封装为[images](只有images这个元素) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel( self.network, chunk_sizes=system_configs.chunk_sizes).cuda() # 此处的DataParallel是作者自己写的一个在模块级别实现数据并行的类,该容器通过 将batch size个输入按照chunk_size分配给指定GPU 来并行化数据。 total_params = 0 for params in self.model.parameters(): # 计算参数量 num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) # 打印总的参数量 if system_configs.opt_algo == "adam": # 参数更新策略:adam self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters()) ) # # 只有requires_grad=True的参数需要optimize elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): # 模型传到GPU self.model.cuda() def train_mode(self): # 模型设置为训练模式 self.network.train() def eval_mode(self): self.network.eval() # 模型设置为推理模式 def train(self, xs, ys, **kwargs): xs = [x for x in xs] ys = [y for y in ys] self.optimizer.zero_grad() loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() focal_loss = focal_loss.mean() pull_loss = pull_loss.mean() push_loss = push_loss.mean() regr_loss = regr_loss.mean() loss.backward() self.optimizer.step() return loss, focal_loss, pull_loss, push_loss, regr_loss def validate(self, xs, ys, **kwargs): # 验证时不需要反向传播 with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): # 学习率设置 print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): # 加载预训练模型 print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def load_params(self, iteration): # 模型加载参数 cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): # 保存模型参数 cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, db, cuda_flag): super(NetworkFactory, self).__init__() # print("[NetworkFactory __init__] db", db ) module_file = "models.{}".format(system_configs.snapshot_name) # print("[NetworkFactory __init__] module_file: {}".format(module_file)) # [NetworkFactory __init__] module_file: models.medical_ExtremeNet nnet_module = importlib.import_module(module_file) # print("[NetworkFactory __init__] nnet_module", nnet_module) self.model = DummyModule(nnet_module.model(db)) self.loss = nnet_module.loss # yezheng: this is last line in models/ExtremeNet.py self.network = Network(self.model, self.loss) self.cuda_flag = cuda_flag if self.cuda_flag: self.network = DataParallel(self.network, chunk_sizes=system_configs.chunk_sizes) total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters()) ) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD( filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001 ) else: raise ValueError("unknown optimizer") # print("[NetworkFactory] system_configs.snapshot_file",system_configs.snapshot_file) # ystem_configs.snapshot_file ./cache/nnet/medical_ExtremeNet/medical_ExtremeNet_{}.pkl def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): if torch.cuda.is_available() and self.cuda_flag: xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] self.optimizer.zero_grad() loss = self.network(xs, ys) loss = loss.mean() loss.backward() self.optimizer.step() return loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): if torch.cuda.is_available() and self.cuda_flag: xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss = self.network(xs, ys) loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): if torch.cuda.is_available() and self.cuda_flag: xs = [x.cuda(non_blocking=True) for x in xs] # print("[NetworkFactory test] list len(xs)", len(xs), # "xs[0]",xs[0].shape, type(xs[0]) )#, "self.model", self.model) return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: if torch.cuda.is_available() and self.cuda_flag: params = torch.load(f) else: params = torch.load(f, map_location = 'cpu') self.model.load_state_dict(params, strict=False) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) print_log("loading model from {}".format(cache_file), system_configs) with open(cache_file, "rb") as f: if torch.cuda.is_available() and self.cuda_flag: params = torch.load(f) else: params = torch.load(f, map_location = 'cpu') self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, flag=False): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) # print("module_file: {}".format(module_file)) # models.CornerNet nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(flag=flag)) self.loss = nnet_module.loss() self.network = Network(self.model, self.loss) self.network = DataParallel(self.network, chunk_sizes=system_configs.chunk_sizes) self.flag = flag # Count total parameters total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("Total parameters: {}".format(total_params)) # Count MACs when input is 360 x 640 x 3 input_test = torch.randn(1, 3, 360, 640).cuda() input_mask = torch.randn(1, 3, 360, 640).cuda() macs, params, = profile(self.model, inputs=(input_test, input_mask), verbose=False) macs, _ = clever_format([macs, params], "%.3f") print('MACs: {}'.format(macs)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) elif system_configs.opt_algo == 'adamW': self.optimizer = torch.optim.AdamW(filter( lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, weight_decay=1e-4) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, iteration, save, viz_split, xs, ys, **kwargs): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] self.optimizer.zero_grad() loss_kp = self.network(iteration, save, viz_split, xs, ys) loss = loss_kp[0] loss_dict = loss_kp[1:] loss = loss.mean() loss.backward() self.optimizer.step() return loss, loss_dict def validate(self, iteration, save, viz_split, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(iteration, save, viz_split, xs, ys) loss = loss_kp[0] loss_dict = loss_kp[1:] loss = loss.mean() return loss, loss_dict def test(self, xs, **kwargs): with torch.no_grad(): # xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def load_params(self, iteration, is_bbox_only=False): cache_file = system_configs.snapshot_file.format(iteration) with open(cache_file, "rb") as f: params = torch.load(f) model_dict = self.model.state_dict() if len(params) != len(model_dict): pretrained_dict = { k: v for k, v in params.items() if k in model_dict } else: pretrained_dict = params model_dict.update(pretrained_dict) self.model.load_state_dict(model_dict) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, db, num_classes): super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) self.model = DummyModule(nnet_module.model(db, num_classes)) self.loss = nnet_module.loss self.network = Network(self.model, self.loss) self.network = DataParallel( self.network, chunk_sizes=system_configs.chunk_sizes).cuda() total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters())) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [x for x in xs] ys = [y for y in ys] self.optimizer.zero_grad() loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() focal_loss = focal_loss.mean() pull_loss = pull_loss.mean() push_loss = push_loss.mean() regr_loss = regr_loss.mean() loss.backward() self.optimizer.step() return loss, focal_loss, pull_loss, push_loss, regr_loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) # self.model.load_state_dict(params) # new addded -- skip mismatched layers new = list(params.items()) pretrained_kvpair = self.model.state_dict() # print(pretrained_kvpair.keys()) for count, (key, value) in enumerate(pretrained_kvpair.items()): try: # print("==>before:", self.model.state_dict()[key]) layer_name, weights = new[count] pretrained_kvpair[key] = weights # print("==>after:", self.model.state_dict()[key]) except Exception as e: print("!!==>", e) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)
class NetworkFactory(object): def __init__(self, db): # db is a MSCOCO instance. # and images, tl, br, center heatmaps and tl,br,center regression are all # prepared into torch.Tensor super(NetworkFactory, self).__init__() module_file = "models.{}".format(system_configs.snapshot_name) # snapshot_name check CenterNet-104.json # snapshot = 5000 but snapshot_name wasn't in CenterNet-104.json # snapshot_name in config.py is None # snapshot_name is set to "CenterNet-104" in train.py # so here module_file = models.CenterNet-104 # it is to say we should go to models/CenterNet-104.py to see the classes print("module_file: {}".format(module_file)) nnet_module = importlib.import_module(module_file) # this is importing the models/CenterNet-104.py # include five functions and 6 classes self.model = DummyModule(nnet_module.model(db)) # nnet_module.model(db) means CenterNet-104.py class model() # model inherit kp # self.loss = nnet_module.loss # CenterNet-104.py loss # loss = AELoss(pull_weight=1e-1, push_weight=1e-1, focal_loss=_neg_loss) # AELoss in kp.py self.network = Network(self.model, self.loss) self.network = DataParallel(self.network, chunk_sizes=system_configs.chunk_sizes).cuda() total_params = 0 for params in self.model.parameters(): num_params = 1 for x in params.size(): num_params *= x total_params += num_params print("total parameters: {}".format(total_params)) if system_configs.opt_algo == "adam": self.optimizer = torch.optim.Adam( filter(lambda p: p.requires_grad, self.model.parameters()) ) elif system_configs.opt_algo == "sgd": self.optimizer = torch.optim.SGD( filter(lambda p: p.requires_grad, self.model.parameters()), lr=system_configs.learning_rate, momentum=0.9, weight_decay=0.0001 ) else: raise ValueError("unknown optimizer") def cuda(self): self.model.cuda() def train_mode(self): self.network.train() def eval_mode(self): self.network.eval() def train(self, xs, ys, **kwargs): xs = [x for x in xs] ys = [y for y in ys] self.optimizer.zero_grad() loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() focal_loss = focal_loss.mean() pull_loss = pull_loss.mean() push_loss = push_loss.mean() regr_loss = regr_loss.mean() loss.backward() self.optimizer.step() return loss, focal_loss, pull_loss, push_loss, regr_loss def validate(self, xs, ys, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] ys = [y.cuda(non_blocking=True) for y in ys] loss_kp = self.network(xs, ys) loss = loss_kp[0] focal_loss = loss_kp[1] pull_loss = loss_kp[2] push_loss = loss_kp[3] regr_loss = loss_kp[4] loss = loss.mean() return loss def test(self, xs, **kwargs): with torch.no_grad(): xs = [x.cuda(non_blocking=True) for x in xs] return self.model(*xs, **kwargs) def set_lr(self, lr): print("setting learning rate to: {}".format(lr)) for param_group in self.optimizer.param_groups: param_group["lr"] = lr def load_pretrained_params(self, pretrained_model): print("loading from {}".format(pretrained_model)) with open(pretrained_model, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def load_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("loading model from {}".format(cache_file)) with open(cache_file, "rb") as f: params = torch.load(f) self.model.load_state_dict(params) def save_params(self, iteration): cache_file = system_configs.snapshot_file.format(iteration) print("saving model to {}".format(cache_file)) with open(cache_file, "wb") as f: params = self.model.state_dict() torch.save(params, f)