def main():
args = make_args()
config = configparser.ConfigParser()
utils.load_config(config, args.config)
for cmd in args.modify:
utils.modify_config(config, cmd)
with open(os.path.expanduser(os.path.expandvars(args.logging)), 'r') as f:
logging.config.dictConfig(yaml.load(f))
height, width = tuple(map(int, config.get('image', 'size').split()))
model_dir = utils.get_model_dir(config)
_, num_parts = utils.get_dataset_mappers(config)
limbs_index = utils.get_limbs_index(config)
path, step, epoch = utils.train.load_model(model_dir)
state_dict = torch.load(path, map_location=lambda storage, loc: storage)
config_channels_dnn = model.ConfigChannels(config, state_dict['dnn'])
dnn = utils.parse_attr(config.get('model', 'dnn'))(config_channels_dnn)
config_channels_stages = model.ConfigChannels(config, state_dict['stages'], config_channels_dnn.channels)
channel_dict = model.channel_dict(num_parts, len(limbs_index))
stages = nn.Sequential(*[utils.parse_attr(s)(config_channels_stages, channel_dict, config_channels_dnn.channels, str(i)) for i, s in enumerate(config.get('model', 'stages').split())])
dnn.load_state_dict(config_channels_dnn.state_dict)
stages.load_state_dict(config_channels_stages.state_dict)
inference = model.Inference(config, dnn, stages)
inference.eval()
logging.info(humanize.naturalsize(sum(var.cpu().numpy().nbytes for var in inference.state_dict().values())))
image = torch.autograd.Variable(torch.randn(args.batch_size, 3, height, width), volatile=True)
path = model_dir + '.onnx'
logging.info('save ' + path)
forward = inference.forward
inference.forward = lambda self, *x: [[output[name] for name in 'parts, limbs'.split(', ')] for output in forward(self, *x)]
torch.onnx.export(inference, image, path, export_params=True, verbose=args.verbose)
def load(self):
try:
path, step, epoch = utils.train.load_model(self.model_dir)
state_dict = torch.load(path,
map_location=lambda storage, loc: storage)
except (FileNotFoundError, ValueError):
step, epoch = 0, 0
state_dict = {name: None for name in ('dnn', 'stages')}
config_channels_dnn = model.ConfigChannels(self.config,
state_dict['dnn'])
dnn = utils.parse_attr(self.config.get('model',
'dnn'))(config_channels_dnn)
config_channels_stages = model.ConfigChannels(
self.config, state_dict['stages'], config_channels_dnn.channels)
channel_dict = model.channel_dict(self.num_parts,
len(self.limbs_index))
stages = nn.Sequential(*[
utils.parse_attr(s)(config_channels_stages, channel_dict,
config_channels_dnn.channels, str(i))
for i, s in enumerate(self.config.get('model', 'stages').split())
])
if config_channels_dnn.state_dict is not None:
dnn.load_state_dict(config_channels_dnn.state_dict)
if config_channels_stages.state_dict is not None:
stages.load_state_dict(config_channels_stages.state_dict)
return step, epoch, dnn, stages
def load(self):
path, step, epoch = utils.train.load_model(self.model_dir)
state_dict = torch.load(path,
map_location=lambda storage, loc: storage)
config_channels_dnn = model.ConfigChannels(self.config,
state_dict['dnn'])
dnn = utils.parse_attr(self.config.get('model',
'dnn'))(config_channels_dnn)
config_channels_stages = model.ConfigChannels(
self.config, state_dict['stages'], config_channels_dnn.channels)
channel_dict = model.channel_dict(self.num_parts,
len(self.limbs_index))
stages = nn.Sequential(*[
utils.parse_attr(s)(config_channels_stages, channel_dict,
config_channels_dnn.channels, str(i))
for i, s in enumerate(self.config.get('model', 'stages').split())
])
dnn.load_state_dict(config_channels_dnn.state_dict)
stages.load_state_dict(config_channels_stages.state_dict)
return step, epoch, dnn, stages
def main():
args = make_args()
config = configparser.ConfigParser()
utils.load_config(config, args.config)
for cmd in args.modify:
utils.modify_config(config, cmd)
with open(os.path.expanduser(os.path.expandvars(args.logging)), 'r') as f:
logging.config.dictConfig(yaml.load(f))
torch.manual_seed(args.seed)
mapper = load_mapper(os.path.expandvars(os.path.expanduser(args.mapper)))
model_dir = utils.get_model_dir(config)
_, num_parts = utils.get_dataset_mappers(config)
limbs_index = utils.get_limbs_index(config)
height, width = tuple(map(int, config.get('image', 'size').split()))
tensor = torch.randn(args.batch_size, 3, height, width)
# PyTorch
try:
path, step, epoch = utils.train.load_model(model_dir)
state_dict = torch.load(path,
map_location=lambda storage, loc: storage)
except (FileNotFoundError, ValueError):
state_dict = {name: None for name in ('dnn', 'stages')}
config_channels_dnn = model.ConfigChannels(config, state_dict['dnn'])
dnn = utils.parse_attr(config.get('model', 'dnn'))(config_channels_dnn)
config_channels_stages = model.ConfigChannels(config, state_dict['stages'],
config_channels_dnn.channels)
channel_dict = model.channel_dict(num_parts, len(limbs_index))
stages = nn.Sequential(*[
utils.parse_attr(s)(config_channels_stages, channel_dict,
config_channels_dnn.channels, str(i))
for i, s in enumerate(config.get('model', 'stages').split())
])
inference = model.Inference(config, dnn, stages)
inference.eval()
state_dict = inference.state_dict()
# TensorFlow
with open(os.path.expanduser(os.path.expandvars(args.path)), 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
image = ops.convert_to_tensor(np.transpose(tensor.cpu().numpy(),
[0, 2, 3, 1]),
name='image')
tf.import_graph_def(graph_def, input_map={'image:0': image})
saver = utils.train.Saver(model_dir, config.getint('save', 'keep'))
with tf.Session(config=tf.ConfigProto(device_count={
'CPU': 1,
'GPU': 0
},
allow_soft_placement=True,
log_device_placement=False)) as sess:
try:
for dst in state_dict:
src, converter = mapper[dst]
if src.isdigit():
state_dict[dst].fill_(float(src))
else:
op = sess.graph.get_operation_by_name(src)
t = op.values()[0]
v = sess.run(t)
state_dict[dst] = torch.from_numpy(converter(v))
val = state_dict[dst].numpy()
print('\t'.join(
list(
map(str, (dst, src, val.shape, utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest())))))
inference.load_state_dict(state_dict)
if args.delete:
logging.warning('delete model directory: ' + model_dir)
shutil.rmtree(model_dir, ignore_errors=True)
saver(
dict(
dnn=inference.dnn.state_dict(),
stages=inference.stages.state_dict(),
), 0)
finally:
if args.debug:
for op in sess.graph.get_operations():
if op.values():
logging.info(op.values()[0])
for name in args.debug:
t = sess.graph.get_tensor_by_name(name + ':0')
val = sess.run(t)
val = np.transpose(val, [0, 3, 1, 2])
print('\t'.join(
map(str, [
name,
'x'.join(map(str, val.shape)),
utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest(),
])))
_tensor = torch.autograd.Variable(tensor, volatile=True)
val = dnn(_tensor).data.numpy()
print('\t'.join(
map(str, [
'x'.join(map(str, val.shape)),
utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest(),
])))
for stage, output in enumerate(inference(_tensor)):
for name, feature in output.items():
val = feature.data.numpy()
print('\t'.join(
map(str, [
'stage%d/%s' % (stage, name),
'x'.join(map(str, val.shape)),
utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest(),
])))
forward = inference.forward
inference.forward = lambda self, *x: list(
forward(self, *x)[-1].values())
with SummaryWriter(model_dir) as writer:
writer.add_graph(inference, (_tensor, ))
def main():
args = make_args()
config = configparser.ConfigParser()
utils.load_config(config, args.config)
for cmd in args.modify:
utils.modify_config(config, cmd)
with open(os.path.expanduser(os.path.expandvars(args.logging)), 'r') as f:
logging.config.dictConfig(yaml.load(f))
torch.manual_seed(args.seed)
mapper = load_mapper(os.path.expandvars(os.path.expanduser(args.mapper)))
model_dir = utils.get_model_dir(config)
_, num_parts = utils.get_dataset_mappers(config)
limbs_index = utils.get_limbs_index(config)
height, width = tuple(map(int, config.get('image', 'size').split()))
tensor = torch.randn(args.batch_size, 3, height, width)
# PyTorch
try:
path, step, epoch = utils.train.load_model(model_dir)
state_dict = torch.load(path,
map_location=lambda storage, loc: storage)
except (FileNotFoundError, ValueError):
state_dict = {name: None for name in ('dnn', 'stages')}
config_channels_dnn = model.ConfigChannels(config, state_dict['dnn'])
dnn = utils.parse_attr(config.get('model', 'dnn'))(config_channels_dnn)
config_channels_stages = model.ConfigChannels(config, state_dict['stages'],
config_channels_dnn.channels)
channel_dict = model.channel_dict(num_parts, len(limbs_index))
stages = nn.Sequential(*[
utils.parse_attr(s)(config_channels_stages, channel_dict,
config_channels_dnn.channels, str(i))
for i, s in enumerate(config.get('model', 'stages').split())
])
inference = model.Inference(config, dnn, stages)
inference.eval()
state_dict = inference.state_dict()
# Caffe
net = caffe.Net(os.path.expanduser(os.path.expandvars(args.prototxt)),
os.path.expanduser(os.path.expandvars(args.caffemodel)),
caffe.TEST)
if args.debug:
logging.info('Caffe variables')
for name, blobs in net.params.items():
for i, blob in enumerate(blobs):
val = blob.data
print('\t'.join(
map(str, [
'%s/%d' % (name, i),
'x'.join(map(str, val.shape)),
utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest(),
])))
logging.info('Caffe features')
input = net.blobs[args.input]
input.reshape(*tensor.size())
input.data[...] = tensor.numpy()
net.forward()
for name, blob in net.blobs.items():
val = blob.data
print('\t'.join(
map(str, [
name,
'x'.join(map(str, val.shape)),
utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest(),
])))
# convert
saver = utils.train.Saver(model_dir, config.getint('save', 'keep'))
try:
for dst in state_dict:
src, transform = mapper[dst]
blobs = [b.data for b in net.params[src]]
blob = transform(blobs)
if isinstance(blob, np.ndarray):
state_dict[dst] = torch.from_numpy(blob)
else:
state_dict[dst].fill_(blob)
val = state_dict[dst].numpy()
logging.info('\t'.join(
list(
map(str, (dst, src, val.shape, utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest())))))
inference.load_state_dict(state_dict)
if args.delete:
logging.warning('delete model directory: ' + model_dir)
shutil.rmtree(model_dir, ignore_errors=True)
saver(
dict(
dnn=inference.dnn.state_dict(),
stages=inference.stages.state_dict(),
), 0)
finally:
for stage, output in enumerate(
inference(torch.autograd.Variable(tensor, volatile=True))):
for name, feature in output.items():
val = feature.data.numpy()
print('\t'.join(
map(str, [
'stage%d/%s' % (stage, name),
'x'.join(map(str, val.shape)),
utils.abs_mean(val),
hashlib.md5(val.tostring()).hexdigest(),
])))