def save_simple_points_ply(out_filename, points, text=False): assert points.shape[1] == 3 assert points.ndim == 2 dirname = path.dirname(out_filename) if not path.isdir(dirname): os.makedirs(dirname) log.info('mkdir -p {}'.format(dirname)) xyz_nxyz = points vertex = np.core.records.fromarrays(xyz_nxyz.T, names='x, y, z', formats='f4, f4, f4') el = plyfile.PlyElement.describe(vertex, 'vertex') ply = plyfile.PlyData([el], text=text) ply.write(out_filename) # Replace \r\n with \n. with open(out_filename, 'rb') as f: content = f.read() beginning = content[:128] rest = content[128:] beginning = beginning.replace(b'ply\r\n', b'ply\n') with open(out_filename, 'wb') as f: f.write(beginning + rest) return ply
def run_command(command): log.debug(subprocess.list2cmdline(command)) start_time = time.time() p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE) return_code = p.wait() log.debug('return code: {}'.format(return_code)) elapsed = time.time() - start_time stdout, stderr = p.communicate() stdout = stdout.decode('utf-8') stderr = stderr.decode('utf-8') if p.returncode != 0: exec_summary = textwrap.dedent( textwrap.dedent(""" Command: {} Return code: {} stdout: {} stderr: {} """).format(subprocess.list2cmdline(command), p.returncode, stdout, stderr)) raise RuntimeError(exec_summary) log.info('{0} took {1:.3f} seconds.'.format(path.basename(command[0]), elapsed)) return p, stdout, stderr
def ensure_dir_exists(dirname, log_mkdir=True): dirname = path.realpath(path.expanduser(dirname)) if not path.isdir(dirname): # `exist_ok` in case of race condition. os.makedirs(dirname, exist_ok=True) if log_mkdir: log.info('mkdir -p {}'.format(dirname)) return dirname
def main(): log.info('Getting all filenames.') syndirs = sorted(glob.glob('/data/render_for_cnn/data/syn_images_cropped_bkg_overlaid/*')) random.seed(42) filenames = [] for syndir in syndirs: modeldirs = sorted(glob.glob(path.join(syndir, '*'))) if is_subset: modeldirs = modeldirs[:10] for modeldir in modeldirs: renderings = sorted(glob.glob(path.join(modeldir, '*'))) if is_subset: renderings = renderings[:7] filenames.extend(renderings) log.info('{} files'.format(len(filenames))) data_base_dir = '/data/mvshape' start_time = time.time() random.seed(42) log.info('Processing rgb images.') for i, filename in enumerate(filenames): m = re.search(r'syn_images_cropped_bkg_overlaid/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a', filename) synset = m.group(1) model_name = m.group(2) v = m.group(3) image_num = int(v) vc_rendering_name = '{}_{:04d}'.format(model_name, image_num) out_filename = path.join(data_base_dir, 'out/shapenetcore/single_rgb_128/{}.png'.format(vc_rendering_name)) assert path.isfile(filename) io_utils.ensure_dir_exists(path.dirname(out_filename)) img = io_utils.read_jpg(filename) assert img.shape[0] == img.shape[1] assert img.shape[2] == 3 resize_method = {0: 'bilinear', 1: 'bicubic', 2: 'lanczos'}[random.randint(0, 2)] resized_img = scipy.misc.imresize(img, (128, 128), interp=resize_method) assert resized_img.dtype == np.uint8 # io_utils.save_array_compressed(out_filename, resized_img) scipy.misc.imsave(out_filename, resized_img) if i % 100 == 0: t_elapsed = (time.time() - start_time) t_remaining = (t_elapsed / (i + 1) * (len(filenames) - i)) log.info('Creating examples in db. {} of {}. elapsed: {:.1f} min, remaining: {:.1f} min'.format(i, len(filenames), t_elapsed / 60, t_remaining / 60)) t_elapsed = (time.time() - start_time) log.info('total elapsed: {:.1f} min'.format(t_elapsed / 60))
def print_status(models, num_all_examples, io_block_times: list, accum_scalars): print('') log.info(' {:08d} of {:08d}. io block: {:.5f}s \n{}'.format( models['metadata']['current_epoch_step'], int(num_all_examples / models['metadata']['batch_size']), np.mean(io_block_times), pprint.pformat( torch_utils.recursive_dict_of_lists_mean(accum_scalars)), ))
def save_model(model, save_dir): current_epoch = model['metadata']['current_epoch'] current_epoch_step = model['metadata']['current_epoch_step'] global_step = model['metadata']['global_step'] filename = path.join( save_dir, 'models0_{:05}_{:07}_{:08}.pth'.format(current_epoch, current_epoch_step, global_step)) io_utils.ensure_dir_exists(path.dirname(filename)) with open(filename, 'wb') as f: log.info('Saving.. {}'.format(filename)) torch.save(model, f) log.info('Saved.')
def save_points_ply(out_filename, points, normals, values, confidence, text=False, scale_normals_by_confidence=True): assert points.shape == normals.shape assert points.shape[1] == 3 assert points.ndim == 2 dirname = path.dirname(out_filename) if not path.isdir(dirname): os.makedirs(dirname) log.info('mkdir -p {}'.format(dirname)) if values.ndim == 1: values = values[:, None] if confidence.ndim == 1: confidence = confidence[:, None] # Replace 0 confidence values with a small number to avoid zero-length normals. confidence = confidence.copy() confidence[np.isclose(confidence, 0.0)] = 1e-3 if scale_normals_by_confidence: normals = normals.copy() * confidence xyz_nxyz = np.concatenate([points, normals, values, confidence], axis=1).astype(np.float32) vertex = np.core.records.fromarrays(xyz_nxyz.T, names='x, y, z, nx, ny, nz, value, confidence', formats='f4, f4, f4, f4, f4, f4, f4, f4') el = plyfile.PlyElement.describe(vertex, 'vertex') ply = plyfile.PlyData([el], text=text) ply.write(out_filename) # Replace \r\n with \n. with open(out_filename, 'rb') as f: content = f.read() beginning = content[:128] rest = content[128:] beginning = beginning.replace(b'ply\r\n', b'ply\n') with open(out_filename, 'wb') as f: f.write(beginning + rest) return ply
def load_model(filename, use_cpu=True): assert filename.endswith('.pth'), filename parts = path.basename(filename.split('.')[0]).split('_') current_epoch = int(parts[1]) current_epoch_step = int(parts[2]) global_step = int(parts[3]) model = eval_utils.load_torch_model(filename, use_cpu=use_cpu) log.info('Loading model {}'.format( (current_epoch, current_epoch_step, global_step))) assert 'metadata' in model # if 'metadata' in model: # assert model['metadata']['global_step'] == global_step # assert model['metadata']['current_epoch'] == current_epoch # assert model['metadata']['current_epoch_step'] == current_epoch_step # else: # model['metadata'] = { # 'global_step': global_step, # 'current_epoch': current_epoch, # 'current_epoch_step': current_epoch_step, # } return model
def data_queue_worker(data_loader: dataset.ExampleLoader2, data_queue: queue.Queue, models: dict, discard_uneven_last_batch: bool): """ :param data_loader: Data source object. Implements next() and reset(). :param data_queue: :param models: Only needed to access the metadata. Read-only. :param discard_uneven_last_batch: Should be False for evaluation mode. Often true for training. :return: """ batch_size = models['metadata']['batch_size'] while True: next_batch = data_loader.next() if next_batch is None or (discard_uneven_last_batch and (len(next_batch[0]) != batch_size)): log.info('Resetting data_loader') data_loader.reset() data_queue.put(None, block=True) break if models['metadata']['name'] == 'shapenetcore_rgb_mv6': batch_data_np = prepare_data_rgb_mv(next_batch) elif models['metadata']['name'] == 'shapenetcore_rgb_voxels': batch_data_np = prepare_data_rgb_voxels(next_batch) elif models['metadata']['name'] == 'shrec12_depth_mv6': batch_data_np = prepare_data_depth_mv(next_batch) elif models['metadata']['name'] == 'shrec12_voxels': batch_data_np = prepare_data_depth_voxels(next_batch) else: log.error('ERROR: Not implemented') sys.exit(1) # move to gpu batch_data = torch_utils.recursive_numpy_to_torch(batch_data_np, cuda=True) data_queue.put(batch_data, block=True) log.info('End of data queue')
def open_locked(filename, mode='r+', sleep_seconds=0.2, verbose=True, **kwargs): """ If `filename` does not exist, `mode` cannot be "r" or "r+". Parent directories are created if they do not exist. """ dirname = path.dirname(filename) if 'r' in mode and not path.isfile(filename): raise FileNotFoundError('{} does not exist. `mode` cannot be {}.'.format(filename, mode)) if not path.isdir(dirname): log.info('mkdir -p {}'.format(dirname)) os.makedirs(dirname) with open(filename, mode, **kwargs) as fd: start_time = time.time() try: i = 0 while True: try: fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) break except IOError as e: if e.errno != errno.EAGAIN: raise e else: time.sleep(sleep_seconds) if verbose and i > 0 and i % 10 == 0: log.warning('Waiting for a locked file {}. Elapsed: {:.3g} seconds.'.format(filename, time.time() - start_time)) i += 1 yield fd fd.flush() finally: fcntl.flock(fd, fcntl.LOCK_UN)
def main(): # syn_images_dir = '/data/render_for_cnn/data/syn_images/' # out_dir = '/tmp/hi' syn_images_dir = '/data/syn_images/' out_dir = '/data/syn_images_final/' assert path.isdir(syn_images_dir) bkg_dir = '/home/daeyun/git/RenderForCNN/datasets/sun2012pascalformat/JPEGImages' bkg_file_list = '/home/daeyun/git/RenderForCNN/datasets/sun2012pascalformat/filelist.txt' syn_images_filenames = glob.glob(path.join(syn_images_dir, '**/*.png'), recursive=True) syn_images_filenames = sorted(syn_images_filenames) with open(bkg_file_list, 'r') as f: content = f.read() fnames = content.strip().split() background_filenames = [path.join(bkg_dir, name) for name in fnames] assert path.isfile(syn_images_filenames[0]) assert path.isfile(background_filenames[0]) log.info('{} object images found'.format(len(syn_images_filenames))) log.info('{} background images found'.format(len(background_filenames))) start_time = time.time() all_params = [ { 'object_image_filename': filename, 'bkg_filenames': background_filenames, 'out_image_filename': path.join(out_dir, '/'.join(filename.split('/')[-3:])) } for filename in syn_images_filenames ] pool = mp.Pool(processes=os.cpu_count() * 2) for i, out_filename in enumerate(pool.imap(render_for_cnn_utils.blend_and_resize, all_params, chunksize=200)): if i % 2000 == 0: elapsed = time.time() - start_time remaining = elapsed / (i + 1) * (len(all_params) - (i + 1)) log.info('{} of {}. Elapsed: {} min. Remaining: {} min'.format(i, len(all_params), int(elapsed / 60), int(remaining / 60)))
def truncate_images(syn_images_cropped_exact_dir, out_dir, ignore_overwrite=True): log.info( 'looking for png files in {}'.format(syn_images_cropped_exact_dir)) assert os.path.isdir(syn_images_cropped_exact_dir) filenames = sorted( list(glob.glob(path.join(syn_images_cropped_exact_dir, '*/*/*.png')))) assert len(filenames) > 0 log.info('found {} images'.format(len(filenames))) # synset filtering new_filenames = [] for filename in filenames: m = re.search( r'syn_images_cropped/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a', filename) synset = m.group(1) if synset not in ['02858304', '02876657', '02924116', '04379243']: continue new_filenames.append(filename) filenames = new_filenames assert len(filenames) > 100 print('Number of images to resize:', len(filenames)) start_time = time.time() params = [{ 'filename': filename, 'out_dir': out_dir, 'ignore_overwrite': ignore_overwrite, } for filename in filenames] pool = mp.Pool() for i, out_filename in enumerate( pool.imap(_truncate_images_worker, params, chunksize=200)): if i % 2000 == 0: elapsed = time.time() - start_time remaining = elapsed / (i + 1) * (len(filenames) - (i + 1)) log.info('{} of {}. Elapsed: {} min. Remaining: {} min'.format( i, len(filenames), int(elapsed / 60), int(remaining / 60)))
def main(): syn_images_dir = '/data/mvshape/shapenetcore/single_rgb_128/' shapenetcore_dir = '/data/shapenetcore/ShapeNetCore.v1/' log.info('Getting all filenames.') syndirs = sorted(glob.glob(path.join(syn_images_dir, '*'))) filenames = [] for syndir in syndirs: modeldirs = sorted(glob.glob(path.join(syndir, '*'))) if is_subset: modeldirs = modeldirs[:10] for modeldir in modeldirs: renderings = sorted(glob.glob(path.join(modeldir, '*.png'))) if is_subset: renderings = renderings[:7] filenames.extend(renderings) # random.seed(42) # if not is_subset: # random.shuffle(filenames) # filenames = filenames[:1000000] random.seed(42) log.info('{} files'.format(len(filenames))) # TODO target_dir = '/data/mvshape/database' if is_subset: sqlite_file_path = join(target_dir, 'shapenetcore_subset.sqlite') else: sqlite_file_path = join(target_dir, 'shapenetcore.sqlite') output_cam_distance_from_origin = 2 log.info('Setting up output directory.') # set up debugging directory. if path.isfile(sqlite_file_path): os.remove(sqlite_file_path) io_utils.ensure_dir_exists(target_dir) # used for making sure there is no duplicate. duplicate_name_check_set = set() log.info('Checking for duplicates. And making sure params.txt exists.') for i, filename in enumerate(filenames): m = re.search(r'single_rgb_128/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a', filename) synset = m.group(1) model_name = m.group(2) v = m.group(3) image_num = int(v) vc_rendering_name = '{}_{}_{:04d}'.format(synset, model_name, image_num) if vc_rendering_name in duplicate_name_check_set: print('duplicate found: ', (filename, vc_rendering_name)) duplicate_name_check_set.add(vc_rendering_name) params_filename = join(syn_images_dir, synset, model_name, 'params.txt') assert path.isfile(params_filename) # Create the database dbm.init(sqlite_file_path) with dbm.db.transaction() as txn: log.info('Creating common objects.') make_dataset('shapenetcore') make_rendering_type('rgb') make_rendering_type('depth') make_rendering_type('normal') make_rendering_type('voxels') make_tag('novelview') make_tag('novelmodel') make_tag('novelclass') make_tag('perspective_input') make_tag('orthographic_input') make_tag('perspective_output') make_tag('orthographic_output') make_tag('viewer_centered') make_tag('object_centered') make_tag('real_world') make_split('train') make_split('test') make_split('validation') # Quote from http://shapenet.cs.stanford.edu/shapenet/obj-zip/ShapeNetCore.v1/README.txt # "The OBJ files have been pre-aligned so that the up direction is the +Y axis, and the front is the +X axis. In addition each model is normalized to fit within a unit cube centered at the origin." oc_output_cam = camera.OrthographicCamera.from_Rt( transforms.lookat_matrix(cam_xyz=(0, 0, output_cam_distance_from_origin), obj_xyz=(0, 0, 0), up=(0, 1, 0)), wh=(128, 128), is_world_to_cam=True) db_oc_output_cam = get_db_camera(oc_output_cam, fov=None) # Prepare all category objects. log.info('Preparing categories.') synset_db_category_map = {} for synset, synset_name in synset_name_pairs: db_category_i, _ = dbm.Category.get_or_create(name=synset_name) synset_db_category_map[synset] = db_category_i txn.commit() # Prepare all mesh model objects. # --------------------------------------------- db_object_map = {} # model_name -> {rendering_type_name -> rendering} db_object_centered_renderings = {} log.info('Preparing mesh model objects.') start_time = time.time() count = 0 for i, filename in enumerate(filenames): m = re.search( r'single_rgb_128/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a', filename) synset = m.group(1) model_name = m.group(2) if model_name not in db_object_map: mesh_filename = join(shapenetcore_dir, synset, model_name, 'model.obj') assert path.isfile(mesh_filename) mesh_filename_suffix = join( '/mesh/shapenetcore/v1', '/'.join(mesh_filename.split('/')[-3:])) db_category = synset_db_category_map[synset] # Must be unique. db_object = dbm.Object.create( name=model_name, category=db_category, dataset=datasets['shapenetcore'], num_vertices=0, # Not needed for now. Easy to fill in later. num_faces=0, mesh_filename=mesh_filename_suffix, ) db_object_map[model_name] = db_object oc_rendering_name = '{}_{}'.format(synset, model_name) assert model_name not in db_object_centered_renderings db_object_centered_renderings[model_name] = { 'output_rgb': dbm.ObjectRendering.create( type=rendering_types['rgb'], camera=db_oc_output_cam, object=db_object, # JPG filename='/shapenetcore/mv20_rgb_128/{}.bin'.format( oc_rendering_name), resolution=128, num_channels=3, set_size=20, is_normalized=False, ), 'output_depth': dbm.ObjectRendering.create( type=rendering_types['depth'], camera=db_oc_output_cam, object=db_object, # Since there is only one gt rendering per model, their id is the same as the model name. filename='/shapenetcore/mv20_depth_128/{}.bin'.format( oc_rendering_name), resolution=128, num_channels=1, set_size=20, is_normalized=False, ), 'output_normal': dbm.ObjectRendering.create( type=rendering_types['normal'], camera=db_oc_output_cam, object=db_object, filename='/shapenetcore/mv20_normal_128/{}.bin'.format( oc_rendering_name), resolution=128, num_channels=3, set_size=20, is_normalized=False, ), 'output_voxels': dbm.ObjectRendering.create( type=rendering_types['voxels'], camera=db_oc_output_cam, object=db_object, filename='/shapenetcore/voxels_32/{}.bin'.format( oc_rendering_name), resolution=32, num_channels=1, set_size=1, is_normalized=False, ) } if count % 5000 == 0: txn.commit() t_elapsed = (time.time() - start_time) t_remaining = (t_elapsed / (i + 1) * (len(filenames) - i)) log.info( 'Creating mesh objects in db. {} of {}. elapsed: {:.1f} min, remaining: {:.1f} min' .format(i, len(filenames), t_elapsed / 60, t_remaining / 60)) count += 1 txn.commit() t_elapsed = time.time() - start_time log.info('created {} mesh objects in db. elapsed: {:.1f} min'.format( count, t_elapsed / 60)) start_time = time.time() log.info('Processing rgb images.') for i, filename in enumerate(filenames): m = re.search( r'single_rgb_128/(.*?)/(.*?)/[^_]+?_[^_]+?_v(\d{4,})_a', filename) synset = m.group(1) model_name = m.group(2) v = m.group(3) image_num = int(v) params_filename = join(syn_images_dir, synset, model_name, 'params.txt') assert path.isfile(params_filename) lines = render_for_cnn_utils.read_params_file(params_filename) Rt = render_for_cnn_utils.get_Rt_from_RenderForCNN_parameters( lines[image_num]) # Input and output cameras # ------------------- input_cam = camera.OrthographicCamera.from_Rt(Rt, wh=(128, 128), is_world_to_cam=True) # 49.1343 degrees is the default fov in blender. db_input_cam = get_db_camera(input_cam, fov=49.1343) input_cam_depth_xyz = (input_cam.pos / la.norm(input_cam.pos)) * 1.5 input_cam_depth_Rt = transforms.lookat_matrix( cam_xyz=input_cam_depth_xyz, obj_xyz=(0, 0, 0), up=input_cam.up_vector) input_cam_depth = camera.OrthographicCamera.from_Rt( input_cam_depth_Rt, wh=(128, 128), is_world_to_cam=True) db_input_cam_depth = get_db_camera(input_cam_depth, fov=49.1343) output_cam_xyz = (input_cam.pos / la.norm( input_cam.pos)) * output_cam_distance_from_origin output_Rt = transforms.lookat_matrix(cam_xyz=output_cam_xyz, obj_xyz=(0, 0, 0), up=input_cam.up_vector) vc_output_cam = camera.OrthographicCamera.from_Rt( output_Rt, wh=(128, 128), is_world_to_cam=True) db_vc_output_cam = get_db_camera(vc_output_cam, fov=None) # --- db_object = db_object_map[model_name] vc_rendering_name = '{}_{}_{:04d}'.format(synset, model_name, image_num) # Viewer centered renderings. # -------------------------------- # Input rgb image: db_object_rendering_input_rgb = dbm.ObjectRendering.create( type=rendering_types['rgb'], camera=db_input_cam, object=db_object, # This should already exist. filename='/shapenetcore/single_rgb_128/{}.png'.format( vc_rendering_name), resolution=128, num_channels=1, set_size=1, is_normalized=False, # False for rgb. ) db_object_rendering_input_depth = dbm.ObjectRendering.create( type=rendering_types['depth'], camera=db_input_cam_depth, object=db_object, filename='/shapenetcore/single_depth_128/{}.bin'.format( vc_rendering_name), resolution=128, num_channels=1, set_size=1, is_normalized=True, ) db_object_rendering_vc_output_rgb = dbm.ObjectRendering.create( type=rendering_types['rgb'], camera=db_vc_output_cam, object=db_object, filename='/shapenetcore/mv20_rgb_128/{}.bin'.format( vc_rendering_name), resolution=128, num_channels=3, set_size=20, is_normalized=False, ) db_object_rendering_vc_output_depth = dbm.ObjectRendering.create( type=rendering_types['depth'], camera=db_vc_output_cam, object=db_object, filename='/shapenetcore/mv20_depth_128/{}.bin'.format( vc_rendering_name), resolution=128, num_channels=1, set_size=20, is_normalized=False, ) db_object_rendering_vc_output_normal = dbm.ObjectRendering.create( type=rendering_types['normal'], camera=db_vc_output_cam, object=db_object, filename='/shapenetcore/mv20_normal_128/{}.bin'.format( vc_rendering_name), resolution=128, num_channels=3, set_size=20, is_normalized=False, ) db_object_rendering_vc_output_voxels = dbm.ObjectRendering.create( type=rendering_types['voxels'], camera=db_vc_output_cam, object=db_object, filename='/shapenetcore/voxels_32/{}.bin'.format( vc_rendering_name), resolution=32, num_channels=1, set_size=1, is_normalized=False, ) # Examples # ---------------- # A row in the `Example` table is just an id for many-to-many references. # View centered example_viewer_centered = dbm.Example.create() dbm.ExampleObjectRendering.create( example=example_viewer_centered, rendering=db_object_rendering_input_rgb) dbm.ExampleObjectRendering.create( example=example_viewer_centered, rendering=db_object_rendering_input_depth) dbm.ExampleObjectRendering.create( example=example_viewer_centered, rendering=db_object_rendering_vc_output_depth) dbm.ExampleObjectRendering.create( example=example_viewer_centered, rendering=db_object_rendering_vc_output_normal) dbm.ExampleObjectRendering.create( example=example_viewer_centered, rendering=db_object_rendering_vc_output_rgb) dbm.ExampleObjectRendering.create( example=example_viewer_centered, rendering=db_object_rendering_vc_output_voxels) dbm.ExampleDataset.create(example=example_viewer_centered, dataset=datasets['shapenetcore']) dbm.ExampleSplit.create(example=example_viewer_centered, split=splits['train']) dbm.ExampleTag.create(example=example_viewer_centered, tag=tags['real_world']) dbm.ExampleTag.create(example=example_viewer_centered, tag=tags['viewer_centered']) dbm.ExampleTag.create(example=example_viewer_centered, tag=tags['perspective_input']) dbm.ExampleTag.create(example=example_viewer_centered, tag=tags['orthographic_output']) dbm.ExampleTag.create(example=example_viewer_centered, tag=tags['novelmodel']) # Object centered example_object_centered = dbm.Example.create() dbm.ExampleObjectRendering.create( example=example_object_centered, rendering=db_object_rendering_input_rgb) dbm.ExampleObjectRendering.create( example=example_object_centered, rendering=db_object_rendering_input_depth) dbm.ExampleObjectRendering.create( example=example_object_centered, rendering=db_object_centered_renderings[model_name] ['output_depth']) dbm.ExampleObjectRendering.create( example=example_object_centered, rendering=db_object_centered_renderings[model_name] ['output_normal']) dbm.ExampleObjectRendering.create( example=example_object_centered, rendering=db_object_centered_renderings[model_name] ['output_rgb']) dbm.ExampleObjectRendering.create( example=example_object_centered, rendering=db_object_centered_renderings[model_name] ['output_voxels']) dbm.ExampleDataset.create(example=example_object_centered, dataset=datasets['shapenetcore']) dbm.ExampleSplit.create(example=example_object_centered, split=splits['train']) dbm.ExampleTag.create(example=example_object_centered, tag=tags['real_world']) dbm.ExampleTag.create(example=example_object_centered, tag=tags['object_centered']) dbm.ExampleTag.create(example=example_object_centered, tag=tags['perspective_input']) dbm.ExampleTag.create(example=example_object_centered, tag=tags['orthographic_output']) dbm.ExampleTag.create(example=example_object_centered, tag=tags['novelmodel']) if i % 5000 == 0: txn.commit() t_elapsed = (time.time() - start_time) t_remaining = (t_elapsed / (i + 1) * (len(filenames) - i)) log.info( 'Creating examples in db. {} of {}. elapsed: {:.1f} min, remaining: {:.1f} min' .format(i, len(filenames), t_elapsed / 60, t_remaining / 60)) txn.commit() dbm.db.commit() t_elapsed = (time.time() - start_time) log.info('total elapsed: {:.1f} min'.format(t_elapsed / 60))
def main(): experiment_name = args.experiment.strip() save_name = args.save_name model_name = args.model_name resume_filename = path.expanduser(args.resume) print_every = args.print_every # sanity checks assert experiment_name in ('opo', 'vpo'), experiment_name assert save_name, save_name assert model_name in encoderdecoder.model_names, model_name assert print_every > 0 learning_rate = get_learning_rate(model_name) log.info(args) if resume_filename: assert resume_filename.endswith('.pth') assert path.isfile(resume_filename) assert '/' + experiment_name in resume_filename # sanity check. doesnt necessarily have to be true # e.g. '/data/mvshape/out/pytorch/shapenetcore_rgb_mv6/vpo/0/models0_00000_0017500_00017500.pth' models = load_model(resume_filename, model_name) else: models = build_model(model_name) # models['metadata']['batch_size'] batch_size = models['metadata']['batch_size'] # output of model is # ((b, 1, 128, 128), (b, 1, 128, 128)) # TODO: this should probably be in gpu helper_torch_modules = encoderdecoder.build_helper_torch_modules() all_params = [] torch_utils.recursive_module_apply( models, lambda m: all_params.extend(m.parameters())) assert len(all_params) > 0 optimizer = optim.Adam(all_params, lr=learning_rate) log.info('Model is ready. Loading training metadata.') # Initialize dataset. ##################################### loader = select_data_loader(model_name=model_name, experiment_name=experiment_name, batch_size=batch_size) log.info('Data is ready.') save_dir = path.join( dataset.base, 'out/pytorch/{}/{}/{}/'.format(model_name, experiment_name, save_name)) # Main loop while True: encoderdecoder.train_epoch( models=models, data_loader=loader, optimizer=optimizer, helper_torch_modules=helper_torch_modules, save_every=2500, print_every=print_every, save_dir=save_dir, )
def train_epoch(models, data_loader: dataset.ExampleLoader2, optimizer: torch.optim.Optimizer, helper_torch_modules: dict, save_every=2500, print_every=50, save_dir=None): log.info('--- Started training. \n{}'.format( pprint.pformat(models['metadata']))) torch_utils.recursive_train_setter(models, is_training=True) gc.collect() batch_size = models['metadata']['batch_size'] num_all_examples = data_loader.total_num_examples assert batch_size == data_loader.batch_size, (batch_size, data_loader.batch_size) # sanity check. not really necessary. just making sure .reset() was called. assert data_loader.current_batch_index == 0, data_loader.current_batch_index current_num_examples = 0 all_scalar_outputs = {} data_queue = queue.Queue(maxsize=1) thread = threading.Thread(target=data_queue_worker, args=[data_loader, data_queue, models, True], daemon=True) thread.start() io_block_times = [] model_name = models['metadata']['name'] models['metadata']['current_epoch'] += 1 log.info('\n### New epoch: {}\n'.format( models['metadata']['current_epoch'])) models['metadata']['current_epoch_step'] = 0 while True: start_time = time.time() batch_data = data_queue.get(block=True) io_block_times.append(time.time() - start_time) if models['metadata']['current_epoch_step'] > 0: if models['metadata']['current_epoch_step'] % print_every == 0: print_status(models, num_all_examples, io_block_times, accum_scalars=all_scalar_outputs) if save_dir is not None and models['metadata'][ 'current_epoch_step'] % save_every == 0: # save_model(models, save_dir=path.join(dataset.base, 'out/pytorch/{}/{}/{}/'.format(model_name, experiment_name, save_name))) save_model(models, save_dir=save_dir) # Check if end of epoch. if batch_data is None: print_status(models, num_all_examples, io_block_times, accum_scalars=all_scalar_outputs) if save_dir is not None: save_model(models, save_dir=save_dir) break assert batch_data['in_image'].size(0) == batch_size optimizer.zero_grad() model_output = get_output_from_model( models, batch_data=batch_data, helper_torch_modules=helper_torch_modules, include_output_tensors=False) total_loss = model_output['loss']['total'] # modifies `all_scalar_outputs` in-place torch_utils.recursive_scalar_dict_merge( { 'loss': model_output['loss'], 'metric': model_output['metric'], }, all_scalar_outputs) # calling `.backward()` seems to trigger garbage collection. without it, it runs out of memory during testing. # this doesn't update any parameters, so it's safe to run when `is_training` is false too. total_loss.backward() optimizer.step() # running `to_python_scalar` triggers a sync point. # total_loss_sum += (to_python_scalar(total_loss) * batch_data['in_image'].size(0)) / this_dataset_total_num_examples current_num_examples += batch_size models['metadata']['current_epoch_step'] += 1 models['metadata']['global_step'] += 1 print('.', end='', flush=True) log.info('\nTrained {} examples\n\n'.format(num_all_examples)) return torch_utils.recursive_dict_of_lists_mean(all_scalar_outputs)
def select_data_loader(model_name, experiment_name, batch_size): if model_name == 'shapenetcore_rgb_mv6': if experiment_name == 'vpo': # loader = dataset.ExampleLoader2(path.join(dataset.base, 'out/splits/shapenetcore_examples_{}/subset_examples.cbor'.format(experiment_name)), # tensors_to_read=('input_rgb', 'target_depth'), shuffle=True, batch_size=batch_size) loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shapenetcore_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_rgb', 'target_depth'), shuffle=True, batch_size=batch_size) elif experiment_name == 'opo': loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shapenetcore_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_rgb', 'target_depth'), shuffle=True, batch_size=batch_size) else: raise NotImplementedError() elif model_name == 'shapenetcore_rgb_voxels': if experiment_name == 'vpo': # TODO######################### loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shapenetcore_examples_{}/subset_examples.cbor'. format(experiment_name)), tensors_to_read=('input_rgb', 'target_voxels'), shuffle=True, batch_size=batch_size) elif experiment_name == 'opo': loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shapenetcore_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_rgb', 'target_voxels'), shuffle=True, batch_size=batch_size) else: raise NotImplementedError() elif model_name == 'shrec12_depth_mv6': if experiment_name == 'vpo': loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shrec12_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_depth', 'target_depth'), shuffle=True, batch_size=batch_size) elif experiment_name == 'opo': loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shrec12_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_depth', 'target_depth'), shuffle=True, batch_size=batch_size) else: raise NotImplementedError() elif model_name == 'shrec12_voxels': if experiment_name == 'vpo': loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shrec12_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_depth', 'target_voxels'), shuffle=True, batch_size=batch_size) elif experiment_name == 'opo': loader = dataset.ExampleLoader2(path.join( dataset.base, 'out/splits/shrec12_examples_{}/all_examples.cbor'.format( experiment_name)), tensors_to_read=('input_depth', 'target_voxels'), shuffle=True, batch_size=batch_size) else: raise NotImplementedError() else: raise NotImplementedError() total_num_examples = loader.total_num_examples log.info('Dataset size: {}'.format(total_num_examples)) return loader
l.depth_and_normal_map.restype = ctypes.c_int l.depth_and_normal_map.argtypes = [ ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"), ctypes.c_size_t, ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"), ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"), ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"), ctypes.c_int, ctypes.c_int, ctypes.c_char_p, ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"), ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"), ] lib = l log.info('Loaded shared library {}'.format(lib._name)) def depth_and_normal_map(vertices, eye, center, up, wh): tri = vertices.astype(np.float32) eye = np.array(eye, np.float32) center = np.array(center, np.float32) up = np.array(up, np.float32) rendered_depth = np.zeros([wh[0], wh[1]], dtype=np.float32) rendered_normal = np.zeros([wh[0], wh[1], 3], dtype=np.float32) resources_dir = path.realpath(path.join(proj_root, 'resources')) lib.depth_and_normal_map(tri, tri.size, eye, center, up, 128, 128, ctypes.c_char_p(resources_dir.encode('ascii')), rendered_depth, rendered_normal)