| | |
|---|
| | | |
|---|
| | | flag = obj.load_state_dict(dst_state, strict=False) |
|---|
| | | # print(flag) |
|---|
| | | |
|---|
| | | # def load_pretrained_model( |
|---|
| | | # path: str, |
|---|
| | | # model: torch.nn.Module, |
|---|
| | | # ignore_init_mismatch: bool, |
|---|
| | | # map_location: str = "cpu", |
|---|
| | | # oss_bucket=None, |
|---|
| | | # scope_map=None, |
|---|
| | | # excludes=None, |
|---|
| | | # ): |
|---|
| | | # """Load a model state and set it to the model. |
|---|
| | | # |
|---|
| | | # Args: |
|---|
| | | # init_param: <file_path>:<src_key>:<dst_key>:<exclude_Keys> |
|---|
| | | # |
|---|
| | | # Examples: |
|---|
| | | # |
|---|
| | | # """ |
|---|
| | | # |
|---|
| | | # obj = model |
|---|
| | | # |
|---|
| | | # if oss_bucket is None: |
|---|
| | | # src_state = torch.load(path, map_location=map_location) |
|---|
| | | # else: |
|---|
| | | # buffer = BytesIO(oss_bucket.get_object(path).read()) |
|---|
| | | # src_state = torch.load(buffer, map_location=map_location) |
|---|
| | | # src_state = src_state["model"] if "model" in src_state else src_state |
|---|
| | | # |
|---|
| | | # if excludes is not None: |
|---|
| | | # for e in excludes.split(","): |
|---|
| | | # src_state = {k: v for k, v in src_state.items() if not k.startswith(e)} |
|---|
| | | # |
|---|
| | | # dst_state = obj.state_dict() |
|---|
| | | # src_state = assigment_scope_map(dst_state, src_state, scope_map) |
|---|
| | | # |
|---|
| | | # if ignore_init_mismatch: |
|---|
| | | # src_state = filter_state_dict(dst_state, src_state) |
|---|
| | | # |
|---|
| | | # logging.debug("Loaded src_state keys: {}".format(src_state.keys())) |
|---|
| | | # logging.debug("Loaded dst_state keys: {}".format(dst_state.keys())) |
|---|
| | | # dst_state.update(src_state) |
|---|
| | | # obj.load_state_dict(dst_state, strict=True) |
|---|
