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)