hongshaorou/tdmpc2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix computation of mem requirements

Browse Source
This commit is contained in:
Nicklas Hansen
2023-12-22 13:57:01 -08:00
parent 34ea3662cd
commit 2929cfdb44
2 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tdmpc2/common/buffer.py
View File

@@ -53,10 +53,10 @@ class Buffer():
bytes_per_ep = sum([ bytes_per_ep = sum([
(v.numel()*v.element_size() if not isinstance(v, TensorDict) \ (v.numel()*v.element_size() if not isinstance(v, TensorDict) \
else sum([x.numel()*x.element_size() for x in v.values()])) \ else sum([x.numel()*x.element_size() for x in v.values()])) \
for k,v in tds.items() for v in tds.values()
]) ])
print(f'Bytes per episode: {bytes_per_ep:,}') print(f'Bytes per episode: {bytes_per_ep:,}')
total_bytes = bytes_per_ep*self._capacity total_bytes = bytes_per_ep * (self._capacity // self.cfg.episode_length)
print(f'Storage required: {total_bytes/1e9:.2f} GB') print(f'Storage required: {total_bytes/1e9:.2f} GB')
# Heuristic: decide whether to use CUDA or CPU memory # Heuristic: decide whether to use CUDA or CPU memory
if 2.5*total_bytes > mem_free: # Insufficient CUDA memory if 2.5*total_bytes > mem_free: # Insufficient CUDA memory

2
tdmpc2/train.py
View File

@@ -1,6 +1,6 @@
import os import os
os.environ['MUJOCO_GL'] = 'egl' os.environ['MUJOCO_GL'] = 'egl'
os.environ['LAZY_LEGACY_OP'] = 0 os.environ['LAZY_LEGACY_OP'] = '0'
import warnings import warnings
warnings.filterwarnings('ignore') warnings.filterwarnings('ignore')
import torch import torch