RuntimeError: CUDA out of memory.

root@Ie640bd7d900d0100b:/hy-tmp/stable-dreamfusion# python main.py --text "a hamburger" --workspace trial -O
Namespace(H=800, O=True, O2=False, W=800, albedo_iters=1000, angle_front=60, angle_overhead=30, backbone='grid', bg_radius=1.4, bound=1, ckpt='latest', cuda_ray=True, density_thresh=10, dir_text=True, dt_gamma=0, eval_interval=10, fovy=60, fovy_range=[40, 70], fp16=True, gui=False, guidance='stable-diffusion', h=64, iters=10000, jitter_pose=False, lambda_entropy=0.0001, lambda_opacity=0, lambda_orient=0.01, light_phi=0, light_theta=60, lr=0.001, max_ray_batch=4096, max_spp=1, max_steps=1024, min_near=0.1, num_steps=64, radius=3, radius_range=[1.0, 1.5], save_mesh=False, seed=0, test=False, text='a hamburger', update_extra_interval=16, upsample_steps=64, w=64, workspace='trial')
NeRFNetwork(
  (encoder): GridEncoder: input_dim=3 num_levels=16 level_dim=2 resolution=16 -> 2048 per_level_scale=1.3819 params=(6119864, 2) gridtype=tiled align_corners=False
  (sigma_net): MLP(
    (net): ModuleList(
      (0): Linear(in_features=32, out_features=64, bias=True)
      (1): Linear(in_features=64, out_features=64, bias=True)
      (2): Linear(in_features=64, out_features=4, bias=True)
    )
  )
  (encoder_bg): FreqEncoder: input_dim=3 degree=6 output_dim=39
  (bg_net): MLP(
    (net): ModuleList(
      (0): Linear(in_features=39, out_features=64, bias=True)
      (1): Linear(in_features=64, out_features=3, bias=True)
    )
  )
)
[INFO] loaded hugging face access token from ./TOKEN!
[INFO] loading stable diffusion...
[INFO] loaded stable diffusion!
[INFO] Trainer: df | 2022-10-17_17-51-17 | cuda | fp16 | trial
[INFO] #parameters: 12249015
[INFO] Loading latest checkpoint ...
[WARN] No checkpoint found, model randomly initialized.
==> Start Training trial Epoch 1, lr=0.010000 ...
  0% 0/100 [00:00<?, ?it/s]Traceback (most recent call last):
  File "main.py", line 153, in <module>
    trainer.train(train_loader, valid_loader, max_epoch)
  File "/hy-tmp/stable-dreamfusion/nerf/utils.py", line 468, in train
    self.train_one_epoch(train_loader)
  File "/hy-tmp/stable-dreamfusion/nerf/utils.py", line 688, in train_one_epoch
    pred_rgbs, pred_ws, loss = self.train_step(data)
  File "/hy-tmp/stable-dreamfusion/nerf/utils.py", line 365, in train_step
    loss = self.guidance.train_step(text_z, pred_rgb)
  File "/hy-tmp/stable-dreamfusion/nerf/sd.py", line 90, in train_step
    noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/diffusers/models/unet_2d_condition.py", line 296, in forward
    sample, res_samples = downsample_block(
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/diffusers/models/unet_blocks.py", line 563, in forward
    hidden_states = attn(hidden_states, context=encoder_hidden_states)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/diffusers/models/attention.py", line 162, in forward
    hidden_states = block(hidden_states, context=context)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/diffusers/models/attention.py", line 211, in forward
    hidden_states = self.attn1(self.norm1(hidden_states)) + hidden_states
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/diffusers/models/attention.py", line 283, in forward
    hidden_states = self._attention(query, key, value)
  File "/usr/local/lib/python3.8/dist-packages/diffusers/models/attention.py", line 292, in _attention
    attention_probs = attention_scores.softmax(dim=-1)
RuntimeError: CUDA out of memory. Tried to allocate 1024.00 MiB (GPU 0; 9.78 GiB total capacity; 7.44 GiB already allocated; 644.06 MiB free; 7.50 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
  0% 0/100 [00:01<?, ?it/s]
ashawkey / stable-dreamfusion

RuntimeError: CUDA out of memory. #41
