def __init__(self, args, project_id, model_name=None, num_classes=None, pytorch_model=None):
self.args = args
self.project_id = project_id
self.model_name = model_name
self.num_classes =num_classes
self.data_dir = None
self.lr = None
# prepare cfg for build model
self.cfg = setup(args=args, project_id=project_id, model_name=model_name, num_classes=num_classes)
super(Detctron2AlObjDetModel, self).__init__(project_id)
self.model, self.device = load_prj_model(project_id=project_id)
if self.model is None:
if pytorch_model:
assert isinstance(
pytorch_model, nn.Module), 'pytorch_model must inherit from torch.nn.Module'
self.model = pytorch_model
print("get a pre-trained model from parameter for project{}".format(project_id))
else:
assert model_name in MODEL_NAME.keys(
), 'model_name must be one of {}'.format(MODEL_NAME.keys())
if not num_classes:
raise ValueError(
"Deep model of project {} is not initialized, please specify the model name and number of classes.".format(
project_id))
self.model = LiuyTrainer.build_model(self.cfg)
self.model = self.model.to(self.device)
print("Initialize a pre-trained model for project{}".format(project_id))
else:
print("load project {} model from file".format(project_id))
print(self.model)
def __init__(self, args, project_id, data_dir):
self.args = args
super(InsSegModel, self).__init__(project_id, data_dir)
# tow ways to get model
# 1：load the model which has been trained
# 2：use the function：LiuyTrainer.build_model(self.cfg)
self.model, self.device = load_prj_model(project_id=project_id)
self.cfg = setup(args=args, project_id=project_id, data_dir=data_dir)
if self.model is None:
self.model = LiuyTrainer.build_model(self.cfg)
self.model = self.model.to(self.device)
print("Initialize a pre-trained model for project{}".format(
project_id))
else:
print("load project {} model from file".format(project_id))
self.trainer = LiuyTrainer(self.cfg, self.model)