Add multi nodes training support

.gitignore

@@ -1 +1,5 @@
+outputs
+cache
 .DS_Store
+*__pycache*
+.idea

README.md

@@ -1,35 +1,8 @@
 ## GET3D: A Generative Model of High Quality 3D Textured Shapes Learned from Images (NeurIPS 2022)<br><sub>Official PyTorch implementation </sub>
 
-![Teaser image](./docs/assets/get3d_model.png)
+This is GET3D implementation with multi-nodes training support.
 
-**GET3D: A Generative Model of High Quality 3D Textured Shapes Learned from Images**<br>
-[Jun Gao](http://www.cs.toronto.edu/~jungao/)
-, [Tianchang Shen](http://www.cs.toronto.edu/~shenti11/)
-, [Zian Wang](http://www.cs.toronto.edu/~zianwang/),
-[Wenzheng Chen](http://www.cs.toronto.edu/~wenzheng/), [Kangxue Yin](https://kangxue.org/)
-, [Daiqing Li](https://scholar.google.ca/citations?user=8q2ISMIAAAAJ&hl=en),
-[Or Litany](https://orlitany.github.io/), [Zan Gojcic](https://zgojcic.github.io/),
-[Sanja Fidler](https://www.cs.toronto.edu/~fidler/) <br>
-**[Paper](https://nv-tlabs.github.io/GET3D/assets/paper.pdf)
-, [Project Page](https://nv-tlabs.github.io/GET3D/)**
-
-Abstract: *As several industries are moving towards modeling massive 3D virtual worlds,
-the need for content creation tools that can scale in terms of the quantity, quality, and
-diversity of 3D content is becoming evident. In our work, we aim to train performant 3D
-generative models that synthesize textured meshes which can be directly consumed by 3D
-rendering engines, thus immediately usable in downstream applications. Prior works on 3D
-generative modeling either lack geometric details, are limited in the mesh topology they
-can produce, typically do not support textures, or utilize neural renderers in the
-synthesis process, which makes their use in common 3D software non-trivial. In this work,
-we introduce GET3D, a Generative model that directly generates Explicit Textured 3D meshes
-with complex topology, rich geometric details, and high fidelity textures. We bridge
-recent success in the differentiable surface modeling, differentiable rendering as well as
-2D Generative Adversarial Networks to train our model from 2D image collections. GET3D is
-able to generate high-quality 3D textured meshes, ranging from cars, chairs, animals,
-motorbikes and human characters to buildings, achieving significant improvements over
-previous methods.*
-
-![Teaser Results](./docs/assets/teaser_result.jpg)
+![Teaser image](./docs/assets/get3d_model.png)
 
 For business inquiries, please visit our website and submit the
 form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/)
@@ -177,4 +150,4 @@ and Daiqing Li and Or Litany and Zan Gojcic and Sanja Fidler},
 booktitle={Advances In Neural Information Processing Systems},
 year={2022}
 }
-```
+```

evaluation_scripts/sample_surface.py

@@ -64,7 +64,7 @@
         [0, options.focal_length_y, options.principal_point_y],
         [0, 0, 1]
     ])
-glctx = dr.RasterizeGLContext()
+glctx = dr.RasterizeCudaContext()
 
 
 def CalcLinearZ(depth):

metrics/metric_utils.py

@@ -46,15 +46,16 @@ def save_image_grid(img, fname, drange, grid_size):
 # ----------------------------------------------------------------------------
 
 class MetricOptions:
-    def __init__(self, G=None, G_kwargs={}, dataset_kwargs={}, num_gpus=1, rank=0, device=None, progress=None, cache=True):
+    def __init__(self, G=None, G_kwargs={}, dataset_kwargs={}, num_gpus=1, rank=0, device=None, progress=None, cache=True, batch_size=None):
         assert 0 <= rank < num_gpus
         self.G = G
         self.G_kwargs = dnnlib.EasyDict(G_kwargs)
         self.dataset_kwargs = dnnlib.EasyDict(dataset_kwargs)
         self.num_gpus = num_gpus
         self.rank = rank
-        self.device = device if device is not None else torch.device('cuda', rank)
+        self.device = device if device is not None else torch.device('cuda')
         self.progress = progress.sub() if progress is not None and rank == 0 else ProgressMonitor()
+        self.batch_size = batch_size
         self.cache = cache
 
 
@@ -227,7 +228,9 @@ def sub(self, tag=None, num_items=None, flush_interval=1000, rel_lo=0, rel_hi=1)
 
 
 # ----------------------------------------------------------------------------
-def compute_feature_stats_for_dataset(opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size=64, data_loader_kwargs=None, max_items=None, **stats_kwargs):
+def compute_feature_stats_for_dataset(opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size=None, data_loader_kwargs=None, max_items=None, **stats_kwargs):
+    if not batch_size:
+        batch_size = opts.batch_size
     dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
     if data_loader_kwargs is None:
         data_loader_kwargs = dict(pin_memory=True, num_workers=3, prefetch_factor=2)

train_3d.py

@@ -21,23 +21,32 @@
 
 
 # ----------------------------------------------------------------------------
-def subprocess_fn(rank, c, temp_dir):
+def subprocess_fn(local_rank, c, temp_dir):
     dnnlib.util.Logger(file_name=os.path.join(c.run_dir, 'log.txt'), file_mode='a', should_flush=True)
 
+    rank = 0
     # Init torch.distributed.
     if c.num_gpus > 1:
+        rank = int(os.environ['RANK'])  # 当前机器编号
+        gpus = torch.cuda.device_count()  # 每台机器的GPU个数
+        rank = rank * gpus + local_rank
+        hosts = int(os.environ['WORLD_SIZE'])  # 机器个数
+        c.num_gpus = hosts * gpus
+
         init_file = os.path.abspath(os.path.join(temp_dir, '.torch_distributed_init'))
         if os.name == 'nt':
             init_method = 'file:///' + init_file.replace('\\', '/')
             torch.distributed.init_process_group(
                 backend='gloo', icfgnit_method=init_method, rank=rank, world_size=c.num_gpus)
         else:
-            init_method = f'file://{init_file}'
+            ip = os.environ.get('MASTER_ADDR', 'localhost')
+            port = os.environ['MASTER_PORT']
+            init_method = f'tcp://{ip}:{port}'
             torch.distributed.init_process_group(
                 backend='nccl', init_method=init_method, rank=rank, world_size=c.num_gpus)
-
+            torch.cuda.set_device(local_rank)
     # Init torch_utils.
-    sync_device = torch.device('cuda', rank) if c.num_gpus > 1 else None
+    sync_device = torch.device('cuda') if c.num_gpus > 1 else None
     training_stats.init_multiprocessing(rank=rank, sync_device=sync_device)
     if rank != 0:
         custom_ops.verbosity = 'none'
@@ -100,9 +109,10 @@ def launch_training(c, desc, outdir, dry_run):
     torch.multiprocessing.set_start_method('spawn', force=True)
     with tempfile.TemporaryDirectory() as temp_dir:
         if c.num_gpus == 1:
-            subprocess_fn(rank=0, c=c, temp_dir=temp_dir)
+            subprocess_fn(local_rank=0, c=c, temp_dir=temp_dir)
         else:
-            torch.multiprocessing.spawn(fn=subprocess_fn, args=(c, temp_dir), nprocs=c.num_gpus)
+            ngpus = torch.cuda.device_count()
+            torch.multiprocessing.spawn(fn=subprocess_fn, args=(c, temp_dir), nprocs=ngpus)
 
 
 # ----------------------------------------------------------------------------
@@ -289,7 +299,8 @@ def main(**kwargs):
     c.image_snapshot_ticks = c.network_snapshot_ticks = opts.snap
     c.random_seed = c.training_set_kwargs.random_seed = opts.seed
     c.data_loader_kwargs.num_workers = opts.workers
-    c.network_snapshot_ticks = 200
+    if opts.gpus <= 8:
+        c.network_snapshot_ticks = 200
     # Sanity checks.
     if c.batch_size % c.num_gpus != 0:
         raise click.ClickException('--batch must be a multiple of --gpus')

training/dataset.py

@@ -289,6 +289,7 @@ def __getitem__(self, idx):
                 or self.data_camera_mode == 'shapenet_motorbike' or self.data_camera_mode == 'ts_house' or self.data_camera_mode == 'ts_animal' \
                 :
             ori_img = cv2.imread(fname, cv2.IMREAD_UNCHANGED)
+            assert ori_img is not None, f'{fname}. No such img.'
             img = ori_img[:, :, :3][..., ::-1]
             mask = ori_img[:, :, 3:4]
             condinfo = np.zeros(2)

training/inference_3d.py

@@ -51,7 +51,7 @@ def inference(
     bias_act._init()
     filtered_lrelu._init()
 
-    device = torch.device('cuda', rank)
+    device = torch.device('cuda')
     np.random.seed(random_seed * num_gpus + rank)
     torch.manual_seed(random_seed * num_gpus + rank)
     torch.backends.cudnn.enabled = True

training/networks_get3d.py

@@ -382,7 +382,7 @@ def extract_3d_shape(
         all_gb_pose = []
         all_uv_mask = []
         if self.dmtet_geometry.renderer.ctx is None:
-            self.dmtet_geometry.renderer.ctx = dr.RasterizeGLContext(device=self.device)
+            self.dmtet_geometry.renderer.ctx = dr.RasterizeCudaContext(device=self.device)
         for v, f in zip(mesh_v, mesh_f):
             uvs, mesh_tex_idx, gb_pos, mask = xatlas_uvmap(
                 self.dmtet_geometry.renderer.ctx, v, f, resolution=texture_resolution)

training/training_loop_3d.py

@@ -116,7 +116,7 @@ def training_loop(
     if num_gpus > 1:
         torch.distributed.barrier()
     start_time = time.time()
-    device = torch.device('cuda', rank)
+    device = torch.device('cuda')
     np.random.seed(random_seed * num_gpus + rank)
     torch.manual_seed(random_seed * num_gpus + rank)
     torch.backends.cudnn.enabled = True
@@ -380,7 +380,7 @@ def training_loop(
             snapshot_data = dict(
                 G=G, D=D, G_ema=G_ema)
             for key, value in snapshot_data.items():
-                if isinstance(value, torch.nn.Module) and not isinstance(value, dr.ops.RasterizeGLContext):
+                if isinstance(value, torch.nn.Module) and not isinstance(value, dr.ops.RasterizeCudaContext):
                     if num_gpus > 1:
                         misc.check_ddp_consistency(value, ignore_regex=r'.*\.[^.]+_(avg|ema|ctx)')
                         for param in misc.params_and_buffers(value):
@@ -406,7 +406,9 @@ def training_loop(
                     with torch.no_grad():
                         result_dict = metric_main.calc_metric(
                             metric=metric, G=snapshot_data['G_ema'],
-                            dataset_kwargs=training_set_kwargs, num_gpus=num_gpus, rank=rank, device=device)
+                            dataset_kwargs=training_set_kwargs, num_gpus=num_gpus,
+                            batch_size=batch_size,
+                            rank=rank, device='cuda')
                     if rank == 0:
                         metric_main.report_metric(result_dict, run_dir=run_dir, snapshot_pkl=snapshot_pkl)
                     stats_metrics.update(result_dict.results)

uni_rep/render/neural_render.py

@@ -56,7 +56,7 @@ def render_mesh(
     ):
         assert not hierarchical_mask
         if self.ctx is None:
-            self.ctx = dr.RasterizeGLContext(device=self.device)
+            self.ctx = dr.RasterizeCudaContext(device=self.device)
 
         mtx_in = torch.tensor(camera_mv_bx4x4, dtype=torch.float32, device=device) if not torch.is_tensor(camera_mv_bx4x4) else camera_mv_bx4x4
         v_pos = xfm_points(mesh_v_pos_bxnx3, mtx_in)  # Rotate it to camera coordinates