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 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) config_channels = model.ConfigChannels(self.config, state_dict) except ValueError: step, epoch = 0, 0 config_channels = model.ConfigChannels(self.config) return step, epoch, config_channels
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)) model_dir = utils.get_model_dir(config) category = utils.get_category(config) anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() path, step, epoch = utils.train.load_model(model_dir) state_dict = torch.load(path, map_location=lambda storage, loc: storage) _model = utils.parse_attr(config.get('model', 'dnn')) dnn = _model(model.ConfigChannels(config, state_dict), anchors, len(category)) logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in dnn.state_dict().values()))) dnn.load_state_dict(state_dict) height, width = tuple(map(int, config.get('image', 'size').split())) image = torch.autograd.Variable( torch.randn(args.batch_size, 3, height, width)) output = dnn(image) state_dict = dnn.state_dict() d = utils.dense(state_dict[args.name]) keep = torch.LongTensor(np.argsort(d)[:int(len(d) * args.keep)]) modifier = utils.channel.Modifier( args.name, state_dict, dnn, lambda name, var: var[keep], lambda name, var, mapper: var[mapper(keep, len(d))], debug=args.debug, ) modifier(output.grad_fn) if args.debug: path = modifier.dot.view( '%s.%s.gv' % (os.path.basename(model_dir), os.path.basename(os.path.splitext(__file__)[0])), os.path.dirname(model_dir)) logging.info(path) assert len(keep) == len(state_dict[args.name]) dnn = _model(model.ConfigChannels(config, state_dict), anchors, len(category)) dnn.load_state_dict(state_dict) dnn(image) if not args.debug: torch.save(state_dict, path)
def __init__(self, args, config): self.args = args self.config = config self.model_dir = utils.get_model_dir(config) self.category = utils.get_category(config) self.anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() self.dnn = utils.parse_attr(config.get('model', 'dnn'))( model.ConfigChannels(config), self.anchors, len(self.category)) self.dnn.eval() logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in self.dnn.state_dict().values()))) if torch.cuda.is_available(): self.dnn.cuda() self.height, self.width = tuple( map(int, config.get('image', 'size').split())) output = self.dnn( torch.autograd.Variable(utils.ensure_device( torch.zeros(1, 3, self.height, self.width)), volatile=True)) _, _, self.rows, self.cols = output.size() self.i, self.j = self.rows // 2, self.cols // 2 self.output = output[:, :, self.i, self.j] dataset = Dataset(self.height, self.width) try: workers = self.config.getint('data', 'workers') except configparser.NoOptionError: workers = multiprocessing.cpu_count() self.loader = torch.utils.data.DataLoader( dataset, batch_size=self.args.batch_size, num_workers=workers)
def __init__(self, args, config): self.args = args self.config = config self.cache_dir = utils.get_cache_dir(config) self.model_dir = utils.get_model_dir(config) self.category = utils.get_category(config, self.cache_dir if os.path.exists(self.cache_dir) else None) self.draw_bbox = utils.visualize.DrawBBox(self.category, colors=args.colors, thickness=args.thickness) self.anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() self.height, self.width = tuple(map(int, config.get('image', 'size').split())) self.path, self.step, self.epoch = utils.train.load_model(self.model_dir) state_dict = torch.load(self.path, map_location=lambda storage, loc: storage) self.dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels(config, state_dict), self.anchors, len(self.category)) self.dnn.load_state_dict(state_dict) self.inference = model.Inference(config, self.dnn, self.anchors) self.inference.eval() if torch.cuda.is_available(): self.inference.cuda() logging.info(humanize.naturalsize(sum(var.cpu().numpy().nbytes for var in self.inference.state_dict().values()))) self.create_cap() self.create_cap_size() self.writer = self.create_writer() self.keys = set(args.keys) self.resize = transform.parse_transform(config, config.get('transform', 'resize_test')) self.transform_image = transform.get_transform(config, config.get('transform', 'image_test').split()) self.transform_tensor = transform.get_transform(config, config.get('transform', 'tensor').split())
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) config_channels = model.ConfigChannels(self.config, state_dict) except ValueError: step, epoch = 0, 0 config_channels = model.ConfigChannels(self.config) dnn = utils.parse_attr(self.config.get('model', 'dnn'))(config_channels, self.anchors, len(self.category)) if config_channels.state_dict is not None: dnn.load_state_dict(config_channels.state_dict) return step, epoch, dnn
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)) cache_dir = utils.get_cache_dir(config) model_dir = utils.get_model_dir(config) category = utils.get_category( config, cache_dir if os.path.exists(cache_dir) else None) anchors = utils.get_anchors(config) anchors = torch.from_numpy(anchors).contiguous() path, step, epoch = utils.train.load_model(model_dir) state_dict = torch.load(path, map_location=lambda storage, loc: storage) dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels( config, state_dict), anchors, len(category)) dnn.load_state_dict(state_dict) height, width = tuple(map(int, config.get('image', 'size').split())) resize = transform.parse_transform(config, config.get('transform', 'resize_test')) transform_image = transform.get_transform( config, config.get('transform', 'image_test').split()) transform_tensor = transform.get_transform( config, config.get('transform', 'tensor').split()) # load image image_bgr = cv2.imread('image.jpg') image_resized = resize(image_bgr, height, width) image = transform_image(image_resized) tensor = transform_tensor(image).unsqueeze(0) # Checksum for key, var in dnn.state_dict().items(): a = var.cpu().numpy() print('\t'.join( map(str, [ key, a.shape, utils.abs_mean(a), hashlib.md5(a.tostring()).hexdigest() ]))) output = dnn(torch.autograd.Variable(tensor, volatile=True)).data for key, a in [ ('image_bgr', image_bgr), ('image_resized', image_resized), ('tensor', tensor.cpu().numpy()), ('output', output.cpu().numpy()), ]: print('\t'.join( map(str, [ key, a.shape, utils.abs_mean(a), hashlib.md5(a.tostring()).hexdigest() ])))
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)) cache_dir = utils.get_cache_dir(config) model_dir = utils.get_model_dir(config) category = utils.get_category(config, cache_dir if os.path.exists(cache_dir) else None) anchors = utils.get_anchors(config) anchors = torch.from_numpy(anchors).contiguous() dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels(config), anchors, len(category)) dnn.eval() logging.info(humanize.naturalsize(sum(var.cpu().numpy().nbytes for var in dnn.state_dict().values()))) state_dict = dnn.state_dict() grouped_dict = group_state(state_dict) try: layers = [] with open(os.path.expanduser(os.path.expandvars(args.file)), 'rb') as f: major, minor, revision, seen = struct.unpack('4i', f.read(16)) logging.info('major=%d, minor=%d, revision=%d, seen=%d' % (major, minor, revision, seen)) total = 0 filesize = os.fstat(f.fileno()).st_size for layer in grouped_dict: group = grouped_dict[layer] for suffix in ['conv.bias', 'bn.bias', 'bn.weight', 'bn.running_mean', 'bn.running_var', 'conv.weight']: if suffix in group: var = group[suffix] size = var.size() cnt = np.multiply.reduce(size) total += cnt key = layer + '.' + suffix val = np.array(struct.unpack('%df' % cnt, f.read(cnt * 4)), np.float32) val = np.reshape(val, size) remaining = filesize - f.tell() logging.info('%s.%s: %s=%f (%s), remaining=%d' % (layer, suffix, 'x'.join(list(map(str, size))), utils.abs_mean(val), hashlib.md5(val.tostring()).hexdigest(), remaining)) layers.append([key, torch.from_numpy(val)]) logging.info('%d parameters assigned' % total) layers[-1][1] = transpose_weight(layers[-1][1], len(anchors)) layers[-2][1] = transpose_bias(layers[-2][1], len(anchors)) finally: if remaining > 0: logging.warning('%d bytes remaining' % remaining) state_dict = collections.OrderedDict(layers) if args.delete: logging.warning('delete model directory: ' + model_dir) shutil.rmtree(model_dir, ignore_errors=True) saver = utils.train.Saver(model_dir, config.getint('save', 'keep'), logger=None) path = saver(state_dict, 0, 0) + saver.ext if args.copy is not None: _path = os.path.expandvars(os.path.expanduser(args.copy)) logging.info('copy %s to %s' % (path, _path)) shutil.copy(path, _path)
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 __init__(self, args, config): self.args = args self.config = config self.model_dir = utils.get_model_dir(config) self.cache_dir = utils.get_cache_dir(config) self.category = utils.get_category(config, self.cache_dir) self.draw_bbox = utils.visualize.DrawBBox(self.category) self.loader = self.get_loader() self.anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() self.path, self.step, self.epoch = utils.train.load_model( self.model_dir) state_dict = torch.load(self.path, map_location=lambda storage, loc: storage) dnn = utils.parse_attr(config.get('model', 'dnn'))( model.ConfigChannels(config, state_dict), self.anchors, len(self.category)) dnn.load_state_dict(state_dict) logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in dnn.state_dict().values()))) self.inference = model.Inference(config, dnn, self.anchors) self.inference.eval() if torch.cuda.is_available(): self.inference.cuda() path = self.model_dir + '.ini' if os.path.exists(path): self._config = configparser.ConfigParser() self._config.read(path) else: logging.warning('training config (%s) not found' % path) self.now = datetime.datetime.now() self.mapper = dict([ (inflection.underscore(name), member()) for name, member in inspect.getmembers( importlib.machinery.SourceFileLoader( '', self.config.get('eval', 'mapper')).load_module()) if inspect.isclass(member) ])
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)) model_dir = utils.get_model_dir(config) category = utils.get_category(config) anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() try: path, step, epoch = utils.train.load_model(model_dir) state_dict = torch.load(path, map_location=lambda storage, loc: storage) except (FileNotFoundError, ValueError): logging.warning('model cannot be loaded') state_dict = None dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels( config, state_dict), anchors, len(category)) logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in dnn.state_dict().values()))) if state_dict is not None: dnn.load_state_dict(state_dict) height, width = tuple(map(int, config.get('image', 'size').split())) image = torch.autograd.Variable( torch.randn(args.batch_size, 3, height, width)) output = dnn(image) state_dict = dnn.state_dict() graph = utils.visualize.Graph(config, state_dict) graph(output.grad_fn) diff = [key for key in state_dict if key not in graph.drawn] if diff: logging.warning('variables not shown: ' + str(diff)) path = graph.dot.view( os.path.basename(model_dir) + '.gv', os.path.dirname(model_dir)) logging.info(path)
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)) model_dir = utils.get_model_dir(config) category = utils.get_category(config) anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() path, step, epoch = utils.train.load_model(model_dir) state_dict = torch.load(path, map_location=lambda storage, loc: storage) dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels( config, state_dict), anchors, len(category)) logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in dnn.state_dict().values()))) dnn.load_state_dict(state_dict) height, width = tuple(map(int, config.get('image', 'size').split())) image = torch.autograd.Variable( torch.randn(args.batch_size, 3, height, width)) output = dnn(image) state_dict = dnn.state_dict() closure = utils.walk.Closure(args.name, state_dict, type(dnn).scope, args.debug) closure(output.grad_fn) d = utils.dense(state_dict[args.name]) channels = torch.LongTensor(np.argsort(d)[int(len(d) * args.remove):]) utils.walk.prune(closure, channels) if args.debug: path = closure.dot.view( os.path.basename(model_dir) + '.gv', os.path.dirname(model_dir)) logging.info(path) else: torch.save(state_dict, path)
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)) category = utils.get_category(config) anchors = torch.from_numpy(utils.get_anchors(config)).contiguous() dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels(config), anchors, len(category)) inference = model.Inference(config, dnn, anchors) inference.train() optimizer = eval(config.get('train', 'optimizer'))(filter(lambda p: p.requires_grad, inference.parameters()), args.learning_rate) scheduler = eval(config.get('train', 'scheduler'))(optimizer) for epoch in range(args.epoch): scheduler.step(epoch) lr = scheduler.get_lr() print('\t'.join(map(str, [epoch] + lr)))
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())) cache_dir = utils.get_cache_dir(config) model_dir = utils.get_model_dir(config) category = utils.get_category( config, cache_dir if os.path.exists(cache_dir) else None) anchors = utils.get_anchors(config) anchors = torch.from_numpy(anchors).contiguous() path, step, epoch = utils.train.load_model(model_dir) state_dict = torch.load(path, map_location=lambda storage, loc: storage) dnn = utils.parse_attr(config.get('model', 'dnn'))(model.ConfigChannels( config, state_dict), anchors, len(category)) inference = model.Inference(config, dnn, anchors) inference.eval() logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in inference.state_dict().values()))) dnn.load_state_dict(state_dict) image = torch.autograd.Variable(torch.randn(args.batch_size, 3, height, width), volatile=True) path = model_dir + '.onnx' logging.info('save ' + path) torch.onnx.export(dnn, image, path, export_params=True, verbose=args.verbose) # PyTorch's bug
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)) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') cache_dir = utils.get_cache_dir(config) _, num_parts = utils.get_dataset_mappers(config) limbs_index = utils.get_limbs_index(config) dnn = utils.parse_attr(config.get('model', 'dnn'))( model.ConfigChannels(config)).to(device) logging.info( humanize.naturalsize( sum(var.cpu().numpy().nbytes for var in dnn.state_dict().values()))) size = tuple(map(int, config.get('image', 'size').split())) draw_points = utils.visualize.DrawPoints(limbs_index, colors=config.get( 'draw_points', 'colors').split()) _draw_points = utils.visualize.DrawPoints(limbs_index, thickness=1) draw_bbox = utils.visualize.DrawBBox() paths = [os.path.join(cache_dir, phase + '.pkl') for phase in args.phase] dataset = utils.data.Dataset( config, utils.data.load_pickles(paths), transform=transform.augmentation.get_transform( config, config.get('transform', 'augmentation').split()), shuffle=config.getboolean('data', 'shuffle'), ) logging.info('num_examples=%d' % len(dataset)) try: workers = config.getint('data', 'workers') except configparser.NoOptionError: workers = multiprocessing.cpu_count() collate_fn = utils.data.Collate( config, transform.parse_transform(config, config.get('transform', 'resize_train')), [size], [dnn(torch.randn(1, 3, *size).to(device)).size()[-2:]], maintain=config.getint('data', 'maintain'), transform_image=transform.get_transform( config, config.get('transform', 'image_train').split()), ) loader = torch.utils.data.DataLoader(dataset, shuffle=True, num_workers=workers, collate_fn=collate_fn) for data in loader: path, size, image, mask, keypoints, yx_min, yx_max, parts, limbs, index = ( t.numpy() if hasattr(t, 'numpy') else t for t in (data[key] for key in 'path, size, image, mask, keypoints, yx_min, yx_max, parts, limbs, index' .split(', '))) for path, size, image, mask, keypoints, yx_min, yx_max, parts, limbs, index in zip( *[ path, size, image, mask, keypoints, yx_min, yx_max, parts, limbs, index ]): logging.info(path + ': ' + 'x'.join(map(str, size))) image = utils.visualize.draw_mask(image, mask, 1) size = yx_max - yx_min target = np.logical_and(*[ np.squeeze(a, -1) > 0 for a in np.split(size, size.shape[-1], -1) ]) keypoints, yx_min, yx_max = (a[target] for a in (keypoints, yx_min, yx_max)) for i, points in enumerate(keypoints): if i == index: image = draw_points(image, points) else: image = _draw_points(image, points) image = draw_bbox(image, yx_min.astype(np.int), yx_max.astype(np.int)) dialog = Visualizer('parts', image, parts) dialog.exec() dialog = Visualizer('limbs', image, limbs) dialog.exec()
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(), ])))
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)) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') cache_dir = utils.get_cache_dir(config) _, num_parts = utils.get_dataset_mappers(config) limbs_index = utils.get_limbs_index(config) dnn = utils.parse_attr(config.get('model', 'dnn'))( model.ConfigChannels(config)).to(device) draw_points = utils.visualize.DrawPoints(limbs_index, colors=config.get( 'draw_points', 'colors').split()) _draw_points = utils.visualize.DrawPoints(limbs_index, thickness=1) draw_bbox = utils.visualize.DrawBBox() batch_size = args.rows * args.cols paths = [os.path.join(cache_dir, phase + '.pkl') for phase in args.phase] dataset = utils.data.Dataset( config, utils.data.load_pickles(paths), transform=transform.augmentation.get_transform( config, config.get('transform', 'augmentation').split()), shuffle=config.getboolean('data', 'shuffle'), ) logging.info('num_examples=%d' % len(dataset)) try: workers = config.getint('data', 'workers') except configparser.NoOptionError: workers = multiprocessing.cpu_count() sizes = utils.train.load_sizes(config) feature_sizes = [ dnn(torch.randn(1, 3, *size).to(device)).size()[-2:] for size in sizes ] collate_fn = utils.data.Collate( config, transform.parse_transform(config, config.get('transform', 'resize_train')), sizes, feature_sizes, maintain=config.getint('data', 'maintain'), transform_image=transform.get_transform( config, config.get('transform', 'image_train').split()), ) loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers, collate_fn=collate_fn) for data in loader: path, size, image, mask, keypoints, yx_min, yx_max, index = ( t.numpy() if hasattr(t, 'numpy') else t for t in ( data[key] for key in 'path, size, image, mask, keypoints, yx_min, yx_max, index'. split(', '))) fig, axes = plt.subplots(args.rows, args.cols) axes = axes.flat if batch_size > 1 else [axes] for ax, path, size, image, mask, keypoints, yx_min, yx_max, index in zip( * [axes, path, size, image, mask, keypoints, yx_min, yx_max, index]): logging.info(path + ': ' + 'x'.join(map(str, size))) image = utils.visualize.draw_mask(image, mask, 1) size = yx_max - yx_min target = np.logical_and(*[ np.squeeze(a, -1) > 0 for a in np.split(size, size.shape[-1], -1) ]) keypoints, yx_min, yx_max = (a[target] for a in (keypoints, yx_min, yx_max)) for i, points in enumerate(keypoints): if i == index: image = draw_points(image, points) else: image = _draw_points(image, points) image = draw_bbox(image, yx_min.astype(np.int), yx_max.astype(np.int)) ax.imshow(image) ax.set_xticks([]) ax.set_yticks([]) fig.tight_layout() mng = plt.get_current_fig_manager() mng.resize(*mng.window.maxsize()) plt.show()