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 __init__(self, flag=False):
super(NetworkFactory, self).__init__()
nnet_module = importlib.import_module('models.LSTR')
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=[16])
self.flag = flag
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 __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 __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")
class NetworkFactory(object):
def __init__(self, flag=False):
super(NetworkFactory, self).__init__()
nnet_module = importlib.import_module('models.LSTR')
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=[16])
self.flag = flag
def cuda(self):
self.model.cuda()
def eval_mode(self):
self.network.eval()
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 __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")
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.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).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):
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): # 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):
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):
# 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)
