2019-12-15 05:13:59 +01:00
|
|
|
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
|
|
|
|
#
|
|
|
|
# Licensed under the BSD 3-Clause License (the "License");
|
|
|
|
# you may not use this file except in compliance with the License.
|
|
|
|
# You may obtain a copy of the License at
|
|
|
|
#
|
|
|
|
# https://opensource.org/licenses/BSD-3-Clause
|
|
|
|
#
|
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
# See the License for the specific language governing permissions and
|
|
|
|
# limitations under the License.
|
|
|
|
import argparse
|
|
|
|
import torch
|
|
|
|
|
|
|
|
|
|
|
|
def add_parser_arguments(parser):
|
|
|
|
parser.add_argument(
|
|
|
|
"--checkpoint-path", metavar="<path>", help="checkpoint filename"
|
|
|
|
)
|
|
|
|
parser.add_argument(
|
|
|
|
"--weight-path", metavar="<path>", help="name of file in which to store weights"
|
|
|
|
)
|
2021-11-09 22:42:18 +01:00
|
|
|
parser.add_argument("--ema", action="store_true", default=False)
|
2019-12-15 05:13:59 +01:00
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser(description="PyTorch ImageNet Training")
|
|
|
|
|
|
|
|
add_parser_arguments(parser)
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
2021-11-09 22:42:18 +01:00
|
|
|
checkpoint = torch.load(args.checkpoint_path, map_location=torch.device("cpu"))
|
2019-12-15 05:13:59 +01:00
|
|
|
|
2021-11-09 22:42:18 +01:00
|
|
|
key = "state_dict" if not args.ema else "ema_state_dict"
|
2019-12-15 05:13:59 +01:00
|
|
|
model_state_dict = {
|
2020-06-27 09:32:20 +02:00
|
|
|
k[len("module.") :] if "module." in k else k: v
|
2019-12-15 05:13:59 +01:00
|
|
|
for k, v in checkpoint["state_dict"].items()
|
|
|
|
}
|
2021-04-13 17:00:33 +02:00
|
|
|
print(f"Loaded model, acc : {checkpoint['best_prec1']}")
|
2019-12-15 05:13:59 +01:00
|
|
|
|
2021-11-09 22:42:18 +01:00
|
|
|
torch.save(model_state_dict, args.weight_path)
|