From 52c61f92a73e3e25b553303671f9b40c67c166a9 Mon Sep 17 00:00:00 2001 From: Tail-19 Date: Tue, 24 Dec 2024 14:56:13 +0800 Subject: [PATCH 1/6] before change device to mps_device --- requirements.txt | 2 +- run_nerf.py | 23 ++++++++++++++++------- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/requirements.txt b/requirements.txt index 168d1b855..d94fb06f6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -torch==1.11.0 +# torch==1.11.0 torchvision>=0.9.1 imageio imageio-ffmpeg diff --git a/run_nerf.py b/run_nerf.py index bc270be86..ff90f7bae 100644 --- a/run_nerf.py +++ b/run_nerf.py @@ -19,10 +19,13 @@ from load_LINEMOD import load_LINEMOD_data -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +# device = torch.device("mps" if torch.mps.is_available() else "cpu") +mps_device = torch.device("mps") np.random.seed(0) DEBUG = False +.to(device) + def batchify(fn, chunk): """Constructs a version of 'fn' that applies to smaller batches. @@ -152,8 +155,8 @@ def render_path(render_poses, hwf, K, chunk, render_kwargs, gt_imgs=None, savedi print(i, time.time() - t) t = time.time() rgb, disp, acc, _ = render(H, W, K, chunk=chunk, c2w=c2w[:3,:4], **render_kwargs) - rgbs.append(rgb.cpu().numpy()) - disps.append(disp.cpu().numpy()) + rgbs.append(rgb.mps().numpy()) + disps.append(disp.mps().numpy()) if i==0: print(rgb.shape, disp.shape) @@ -275,7 +278,9 @@ def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=F raw2alpha = lambda raw, dists, act_fn=F.relu: 1.-torch.exp(-act_fn(raw)*dists) dists = z_vals[...,1:] - z_vals[...,:-1] - dists = torch.cat([dists, torch.Tensor([1e10]).expand(dists[...,:1].shape)], -1) # [N_rays, N_samples] + # print(type(dists)) + # print(dists.device) + dists = torch.cat([dists, torch.Tensor([1e10]).expand(dists[...,:1].shape).to("mps")], -1) # [N_rays, N_samples] dists = dists * torch.norm(rays_d[...,None,:], dim=-1) @@ -765,7 +770,8 @@ def train(): img_loss = img2mse(rgb, target_s) trans = extras['raw'][...,-1] loss = img_loss - psnr = mse2psnr(img_loss) + with torch.device(device): + psnr = mse2psnr(img_loss) if 'rgb0' in extras: img_loss0 = img2mse(extras['rgb0'], target_s) @@ -873,6 +879,9 @@ def train(): if __name__=='__main__': - torch.set_default_tensor_type('torch.cuda.FloatTensor') - + # torch.set_default_tensor_type('torch.mps.FloatTensor') + torch.set_default_tensor_type('torch.FloatTensor') + + torch.set_default_dtype(torch.float32) + torch.set_default_device(torch.device("mps")) train() From 5a817b8c0e861ad0b5098976c88d0cf754f52179 Mon Sep 17 00:00:00 2001 From: Tail-19 Date: Tue, 24 Dec 2024 14:58:36 +0800 Subject: [PATCH 2/6] change to mps_device --- run_nerf.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/run_nerf.py b/run_nerf.py index ff90f7bae..9aaaa7161 100644 --- a/run_nerf.py +++ b/run_nerf.py @@ -24,8 +24,6 @@ np.random.seed(0) DEBUG = False -.to(device) - def batchify(fn, chunk): """Constructs a version of 'fn' that applies to smaller batches. @@ -191,14 +189,14 @@ def create_nerf(args): skips = [4] model = NeRF(D=args.netdepth, W=args.netwidth, input_ch=input_ch, output_ch=output_ch, skips=skips, - input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device) + input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(mps_device) grad_vars = list(model.parameters()) model_fine = None if args.N_importance > 0: model_fine = NeRF(D=args.netdepth_fine, W=args.netwidth_fine, input_ch=input_ch, output_ch=output_ch, skips=skips, - input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device) + input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(mps_device) grad_vars += list(model_fine.parameters()) network_query_fn = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn, @@ -653,7 +651,7 @@ def train(): render_kwargs_test.update(bds_dict) # Move testing data to GPU - render_poses = torch.Tensor(render_poses).to(device) + render_poses = torch.Tensor(render_poses).to(mps_device) # Short circuit if only rendering out from trained model if args.render_only: @@ -697,10 +695,10 @@ def train(): # Move training data to GPU if use_batching: - images = torch.Tensor(images).to(device) - poses = torch.Tensor(poses).to(device) + images = torch.Tensor(images).to(mps_device) + poses = torch.Tensor(poses).to(mps_device) if use_batching: - rays_rgb = torch.Tensor(rays_rgb).to(device) + rays_rgb = torch.Tensor(rays_rgb).to(mps_device) N_iters = 200000 + 1 @@ -734,7 +732,7 @@ def train(): # Random from one image img_i = np.random.choice(i_train) target = images[img_i] - target = torch.Tensor(target).to(device) + target = torch.Tensor(target).to(mps_device) pose = poses[img_i, :3,:4] if N_rand is not None: @@ -826,7 +824,7 @@ def train(): os.makedirs(testsavedir, exist_ok=True) print('test poses shape', poses[i_test].shape) with torch.no_grad(): - render_path(torch.Tensor(poses[i_test]).to(device), hwf, K, args.chunk, render_kwargs_test, gt_imgs=images[i_test], savedir=testsavedir) + render_path(torch.Tensor(poses[i_test]).to(mps_device), hwf, K, args.chunk, render_kwargs_test, gt_imgs=images[i_test], savedir=testsavedir) print('Saved test set') From d5dbac89fbdc016ea48c4e61d9871723df73fe0c Mon Sep 17 00:00:00 2001 From: Tail-19 Date: Tue, 24 Dec 2024 15:03:30 +0800 Subject: [PATCH 3/6] run success --- run_nerf.py | 3 +-- run_nerf_helpers.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/run_nerf.py b/run_nerf.py index 9aaaa7161..da2c5b12a 100644 --- a/run_nerf.py +++ b/run_nerf.py @@ -768,8 +768,7 @@ def train(): img_loss = img2mse(rgb, target_s) trans = extras['raw'][...,-1] loss = img_loss - with torch.device(device): - psnr = mse2psnr(img_loss) + psnr = mse2psnr(img_loss) if 'rgb0' in extras: img_loss0 = img2mse(extras['rgb0'], target_s) diff --git a/run_nerf_helpers.py b/run_nerf_helpers.py index bc6ee779d..e7cb73836 100644 --- a/run_nerf_helpers.py +++ b/run_nerf_helpers.py @@ -7,7 +7,7 @@ # Misc img2mse = lambda x, y : torch.mean((x - y) ** 2) -mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.])) +mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]).to(x.device)) to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8) From 91e73ab326b8d15e9c6009cefda3dc491ca78ecf Mon Sep 17 00:00:00 2001 From: Tail-19 Date: Sat, 28 Dec 2024 00:15:50 +0800 Subject: [PATCH 4/6] bug fixed --- run_nerf.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/run_nerf.py b/run_nerf.py index da2c5b12a..e91939109 100644 --- a/run_nerf.py +++ b/run_nerf.py @@ -153,8 +153,8 @@ def render_path(render_poses, hwf, K, chunk, render_kwargs, gt_imgs=None, savedi print(i, time.time() - t) t = time.time() rgb, disp, acc, _ = render(H, W, K, chunk=chunk, c2w=c2w[:3,:4], **render_kwargs) - rgbs.append(rgb.mps().numpy()) - disps.append(disp.mps().numpy()) + rgbs.append(rgb.cpu().numpy()) + disps.append(disp.cpu().numpy()) if i==0: print(rgb.shape, disp.shape) From 5cb19f741d28ab746ac0a4d48ccf8fefd4b0f5a1 Mon Sep 17 00:00:00 2001 From: Tail-19 Date: Mon, 13 Jan 2025 18:25:54 +0800 Subject: [PATCH 5/6] changes for pr --- requirements.txt | 2 +- run_nerf.py | 19 +++++++++---------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/requirements.txt b/requirements.txt index d94fb06f6..f3e3337e3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -# torch==1.11.0 +torch>=1.11.0 torchvision>=0.9.1 imageio imageio-ffmpeg diff --git a/run_nerf.py b/run_nerf.py index e91939109..c81ac3eb4 100644 --- a/run_nerf.py +++ b/run_nerf.py @@ -19,8 +19,7 @@ from load_LINEMOD import load_LINEMOD_data -# device = torch.device("mps" if torch.mps.is_available() else "cpu") -mps_device = torch.device("mps") +device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.mps.device_count() > 0 else "cpu") np.random.seed(0) DEBUG = False @@ -189,14 +188,14 @@ def create_nerf(args): skips = [4] model = NeRF(D=args.netdepth, W=args.netwidth, input_ch=input_ch, output_ch=output_ch, skips=skips, - input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(mps_device) + input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device) grad_vars = list(model.parameters()) model_fine = None if args.N_importance > 0: model_fine = NeRF(D=args.netdepth_fine, W=args.netwidth_fine, input_ch=input_ch, output_ch=output_ch, skips=skips, - input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(mps_device) + input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs).to(device) grad_vars += list(model_fine.parameters()) network_query_fn = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn, @@ -651,7 +650,7 @@ def train(): render_kwargs_test.update(bds_dict) # Move testing data to GPU - render_poses = torch.Tensor(render_poses).to(mps_device) + render_poses = torch.Tensor(render_poses).to(device) # Short circuit if only rendering out from trained model if args.render_only: @@ -695,10 +694,10 @@ def train(): # Move training data to GPU if use_batching: - images = torch.Tensor(images).to(mps_device) - poses = torch.Tensor(poses).to(mps_device) + images = torch.Tensor(images).to(device) + poses = torch.Tensor(poses).to(device) if use_batching: - rays_rgb = torch.Tensor(rays_rgb).to(mps_device) + rays_rgb = torch.Tensor(rays_rgb).to(device) N_iters = 200000 + 1 @@ -732,7 +731,7 @@ def train(): # Random from one image img_i = np.random.choice(i_train) target = images[img_i] - target = torch.Tensor(target).to(mps_device) + target = torch.Tensor(target).to(device) pose = poses[img_i, :3,:4] if N_rand is not None: @@ -823,7 +822,7 @@ def train(): os.makedirs(testsavedir, exist_ok=True) print('test poses shape', poses[i_test].shape) with torch.no_grad(): - render_path(torch.Tensor(poses[i_test]).to(mps_device), hwf, K, args.chunk, render_kwargs_test, gt_imgs=images[i_test], savedir=testsavedir) + render_path(torch.Tensor(poses[i_test]).to(device), hwf, K, args.chunk, render_kwargs_test, gt_imgs=images[i_test], savedir=testsavedir) print('Saved test set') From 10c14923e31b388cb05365aaf7e480fffec3fd7a Mon Sep 17 00:00:00 2001 From: Hanlu Date: Thu, 25 Sep 2025 18:09:42 +0800 Subject: [PATCH 6/6] Update README with fork details for M1 chip Added information about the fork for M1 chip compatibility. --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 27b89ba71..da78091ed 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +This is a fork created by Tail-19, aiming to make this repository work on M1 chip. The project has finished 🎉 + # NeRF-pytorch