def __init__(self):
self.torchDeviceProcess = TorchDeviceProcess()
self.modelFactory = ModelFactory()
self.modelWeightInit = ModelWeightInit()
self.torchDeviceProcess.initTorch()
self.best_value = 0
self.is_multi_gpu = False
class ModelNetShow():
def __init__(self):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.show_process = ModelShow()
def model_show(self, model_path):
input_x = torch.randn(1, 3, 224, 224)
self.show_process.set_input(input_x)
model = self.model_factory.get_model(model_path)
self.show_process.show_from_model(model)
def backbone_show(self, backbone_path):
input_x = torch.randn(1, 3, 224, 224)
self.show_process.set_input(input_x)
model = self.backbone_factory.get_base_model(backbone_path)
self.show_process.show_from_model(model)
def onnx_show(self, onnx_path):
input_x = torch.randn(1, 3, 640, 352)
self.show_process.set_input(input_x)
self.show_process.show_from_onnx(onnx_path)
def pc_model_show(self, backbone_path):
input_x = torch.randn(1, 3, 1024)
self.show_process.set_input(input_x)
model = self.model_factory.get_model(backbone_path)
self.show_process.show_from_model(model)
def pc_backbone_show(self, backbone_path):
input_x = torch.randn(1, 3, 1024)
self.show_process.set_input(input_x)
model = self.backbone_factory.get_base_model(backbone_path)
self.show_process.show_from_model(model)
class ModelConverter():
def __init__(self, input_size=(352, 640)):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.converter = TorchConvertOnnx()
self.input_size = input_size # w * h
def model_convert(self, model_path, weight_path, save_dir):
input_x = torch.randn(1, 3, self.input_size[1], self.input_size[0])
self.converter.set_input(input_x)
self.converter.set_save_dir(save_dir)
model = self.model_factory.get_model(model_path)
self.converter.torch2onnx(model, weight_path)
def base_model_convert(self, base_model_path, weight_path, save_dir):
input_x = torch.randn(1, 3, self.input_size[1], self.input_size[0])
self.converter.set_input(input_x)
self.converter.set_save_dir(save_dir)
model = self.backbone_factory.get_base_model(base_model_path)
self.converter.torch2onnx(model, weight_path)
def model_print(model_name):
model_factory = ModelFactory()
input_x = torch.randn(1, 3, 32, 32)
model = model_factory.get_model(model_name)
model.print_block_name()
class TorchModelProcess():
def __init__(self):
self.torchDeviceProcess = TorchDeviceProcess()
self.modelFactory = ModelFactory()
self.modelWeightInit = ModelWeightInit()
self.torchDeviceProcess.initTorch()
self.best_value = 0
self.is_multi_gpu = False
def initModel(self, cfgPath, gpuId, is_multi_gpu=False):
self.is_multi_gpu = is_multi_gpu
self.torchDeviceProcess.setGpuId(gpuId)
model = self.modelFactory.get_model(cfgPath)
self.modelWeightInit.initWeight(model)
return model
def loadPretainModel(self, weightPath, model):
if weightPath is not None:
if os.path.exists(weightPath):
print("Loading pretainModel from {}".format(weightPath))
model_dict = model.state_dict()
checkpoint = torch.load(weightPath)
pretrained_dict = checkpoint['model']
# pretrained_dict = self.filter_param_dict(pretrained_dict)
new_pretrained_dict = {}
for k, v in pretrained_dict.items():
if k in model_dict and v.shape == model_dict[k].shape:
new_pretrained_dict[k] = v
# print("Load pretrained parameters:")
# for k, v in new_pretrained_dict.items():
# print(k, v.shape)
model_dict.update(new_pretrained_dict)
model.load_state_dict(model_dict)
else:
print("pretain model %s not exist" % weightPath)
def loadLatestModelWeight(self, weightPath, model):
count = self.torchDeviceProcess.getCUDACount()
checkpoint = None
if os.path.exists(weightPath):
if count > 1:
checkpoint = torch.load(weightPath, map_location='cpu')
state = self.convert_state_dict(checkpoint['model'])
model.load_state_dict(state)
else:
checkpoint = torch.load(weightPath, map_location='cpu')
model.load_state_dict(checkpoint['model'])
else:
print("Loading %s fail" % weightPath)
return checkpoint
def saveLatestModel(self,
latestWeightsFile,
model,
optimizer,
epoch=0,
best_value=0):
# Save latest checkpoint
checkpoint = {
'epoch': epoch,
'best_value': best_value,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(checkpoint, latestWeightsFile)
def getLatestModelValue(self, checkpoint):
start_epoch = 0
value = -1
if checkpoint:
if checkpoint.get('epoch') is not None:
start_epoch = checkpoint['epoch'] + 1
if checkpoint.get('best_value') is not None:
value = checkpoint['best_value']
self.setModelBestValue(value)
return start_epoch, value
def setModelBestValue(self, value):
self.best_value = value
def saveBestModel(self, value, latest_weights_file, best_weights_file):
if value >= self.best_value:
self.best_value = value
os.system('cp {} {}'.format(
latest_weights_file,
best_weights_file,
))
return self.best_value
def modelTrainInit(self, model):
count = self.torchDeviceProcess.getCUDACount()
if count > 1 and self.is_multi_gpu:
print('Using ', count, ' GPUs')
model = nn.DataParallel(model)
model = model.to(self.torchDeviceProcess.device)
return model
def modelTestInit(self, model):
model = model.to(self.torchDeviceProcess.device)
return model
def model_clip_grad(self, model):
nn.utils.clip_grad_norm_(model.parameters(), max_norm=20, norm_type=2)
def convert_state_dict(self, state_dict):
"""Converts a state dict saved from a dataParallel module to normal
module state_dict inplace
:param state_dict is the loaded DataParallel model_state
"""
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
return new_state_dict
def getDevice(self):
return self.torchDeviceProcess.device
def filter_param_dict(self,
state_dict: dict,
include: str = None,
exclude: str = None):
assert isinstance(state_dict, dict)
include_re = None
if include is not None:
include_re = re.compile(include)
exclude_re = None
if exclude is not None:
exclude_re = re.compile(exclude)
res_dict = {}
for k, p in state_dict.items():
if include_re is not None:
if include_re.match(k) is None:
continue
if exclude_re is not None:
if exclude_re.match(k) is not None:
continue
res_dict[k] = p
return res_dict
def __init__(self, input_size=(352, 640)):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.converter = TorchConvertOnnx()
self.input_size = input_size # w * h
def __init__(self):
self.backbone_factory = BackboneFactory()
self.model_factory = ModelFactory()
self.show_process = ModelShow()