mirror of https://github.com/JDAI-CV/fast-reid.git
parent
95a3c62ad2
commit
325854c327
|
@ -1,48 +0,0 @@
|
||||||
# encoding: utf-8
|
|
||||||
"""
|
|
||||||
@author: sherlock
|
|
||||||
@contact: sherlockliao01@gmail.com
|
|
||||||
"""
|
|
||||||
|
|
||||||
import sys
|
|
||||||
|
|
||||||
import torch
|
|
||||||
sys.path.append('../..')
|
|
||||||
from fastreid.config import get_cfg
|
|
||||||
from fastreid.engine import default_argument_parser, default_setup
|
|
||||||
from fastreid.modeling.meta_arch import build_model
|
|
||||||
from fastreid.export.tensorflow_export import export_tf_reid_model
|
|
||||||
from fastreid.export.tf_modeling import TfMetaArch
|
|
||||||
|
|
||||||
|
|
||||||
def setup(args):
|
|
||||||
"""
|
|
||||||
Create configs and perform basic setups.
|
|
||||||
"""
|
|
||||||
cfg = get_cfg()
|
|
||||||
# cfg.merge_from_file(args.config_file)
|
|
||||||
cfg.merge_from_list(args.opts)
|
|
||||||
cfg.freeze()
|
|
||||||
default_setup(cfg, args)
|
|
||||||
return cfg
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
args = default_argument_parser().parse_args()
|
|
||||||
print("Command Line Args:", args)
|
|
||||||
cfg = setup(args)
|
|
||||||
cfg.defrost()
|
|
||||||
cfg.MODEL.BACKBONE.NAME = "build_resnet_backbone"
|
|
||||||
cfg.MODEL.BACKBONE.DEPTH = 50
|
|
||||||
cfg.MODEL.BACKBONE.LAST_STRIDE = 1
|
|
||||||
# If use IBN block in backbone
|
|
||||||
cfg.MODEL.BACKBONE.WITH_IBN = False
|
|
||||||
cfg.MODEL.BACKBONE.PRETRAIN = False
|
|
||||||
|
|
||||||
from torchvision.models import resnet50
|
|
||||||
# model = TfMetaArch(cfg)
|
|
||||||
model = resnet50(pretrained=False)
|
|
||||||
# model.load_params_wo_fc(torch.load('logs/bjstation/res50_baseline_v0.4/ckpts/model_epoch80.pth'))
|
|
||||||
model.eval()
|
|
||||||
dummy_inputs = torch.randn(1, 3, 256, 128)
|
|
||||||
export_tf_reid_model(model, dummy_inputs, 'reid_tf.pb')
|
|
Loading…
Reference in New Issue