Hi Danijar, I am currently trying to use higher image resolutions li

I worked a little bit more on this topic and found out that the <code class="notransla

Calling JAXAgent train gets stuck if using larger image sizes (inside Ninjax) about dreamerv3 HOT 3 OPEN

schneimo commented on June 12, 2024

Calling JAXAgent train gets stuck if using larger image sizes (inside Ninjax)

from dreamerv3.

Comments (3)

edwhu commented on June 12, 2024

Sometimes the trace can take a while with old GPUs, I've waited around 10 minutes for a TitanX workstation before.

You can try making the CNN smaller to see if that speeds up compilation time. You can also try incrementally increasing the resolution and check if the trace time increases.

from dreamerv3.

schneimo commented on June 12, 2024

Thanks.

I am not sure if time and compute power is really the problem. Even after 24 hours, it did not trace on an A100. But I will test how tracing time increases with increasing image resolution and report my findings here.

from dreamerv3.

schneimo commented on June 12, 2024

I worked a little bit more on this topic and found out that the train function of class Agent is called completely since when it is decorated with an additional timer, the timer gets executed.

Furthermore, I tracked the problem a little bit more down and it seems to arise in the try block of pure inside the Ninjax module.

dreamerv3/dreamerv3/ninjax.py

Lines 60 to 101 in 8fa35f8

    
           def pure(fun, nested=False): 
        
             """Wrap an impure function that uses global state to explicitly pass the 
        
             state in and out. The result is a pure function that is composable with JAX 
        
             transformation. The pure function can be used as follows: 
        
             `out, state = fun(state, rng, *args, **kwargs)`.""" 
        
             def purified( 
        
                 state, rng, *args, create=None, modify=None, ignore=None, **kwargs): 
        
               context = CONTEXT.get(threading.get_ident(), None) 
        
               if context: 
        
                 create = create if create is not None else context.create 
        
                 modify = modify if modify is not None else context.modify 
        
                 ignore = ignore if ignore is not None else context.ignore 
        
                 assert context.create or not create, 'Parent context disabled create.' 
        
                 assert context.modify or not modify, 'Parent context disabled modify.' 
        
                 assert not context.ignore or ignore, 'Parent context enabled ignore.' 
        
               else: 
        
                 create = create if create is not None else True 
        
                 modify = modify if modify is not None else True 
        
                 ignore = ignore if ignore is not None else False 
        
               if not isinstance(state, dict): 
        
                 raise ValueError('Must provide a dict as state.') 
        
               if context and (not nested): 
        
                 raise RuntimeError( 
        
                     f'You are trying to call pure {fun.__name__}() inside pure ' 
        
                     f'{context.name}(). Is that intentional? If you want to nest pure ' 
        
                     f'functions, use pure(..., nested=True) for the inner function.') 
        
                 # raise RuntimeError( 
        
                 #     f'If you want to nest run() calls, use nested=True. ({context})') 
        
               before = context 
        
               try: 
        
                 name = fun.__name__ 
        
                 if rng.shape == (): 
        
                   rng = jax.random.PRNGKey(rng) 
        
                 context = Context(state.copy(), rng, create, modify, ignore, [], name) 
        
                 CONTEXT[threading.get_ident()] = context 
        
                 out = fun(*args, **kwargs) 
        
                 state = dict(context) 
        
                 return out, state 
        
               finally: 
        
                 CONTEXT[threading.get_ident()] = before 
        
             purified.pure = True 
        
             return purified

from dreamerv3.

Calling JAXAgent train gets stuck if using larger image sizes (inside Ninjax) about dreamerv3 HOT 3 OPEN

Comments (3)

Related Issues (20)

Recommend Projects

React

Vue.js

Typescript

TensorFlow

Django

Laravel

D3

Recommend Topics

javascript

web

server

Machine learning

Visualization

Game

Recommend Org

Facebook

Microsoft

Google

Alibaba

D3

Tencent

	def pure(fun, nested=False):
	"""Wrap an impure function that uses global state to explicitly pass the
	state in and out. The result is a pure function that is composable with JAX
	transformation. The pure function can be used as follows:
	`out, state = fun(state, rng, args, *kwargs)`."""
	def purified(
	state, rng, args, create=None, modify=None, ignore=None, *kwargs):
	context = CONTEXT.get(threading.get_ident(), None)
	if context:
	create = create if create is not None else context.create
	modify = modify if modify is not None else context.modify
	ignore = ignore if ignore is not None else context.ignore
	assert context.create or not create, 'Parent context disabled create.'
	assert context.modify or not modify, 'Parent context disabled modify.'
	assert not context.ignore or ignore, 'Parent context enabled ignore.'
	else:
	create = create if create is not None else True
	modify = modify if modify is not None else True
	ignore = ignore if ignore is not None else False
	if not isinstance(state, dict):
	raise ValueError('Must provide a dict as state.')
	if context and (not nested):
	raise RuntimeError(
	f'You are trying to call pure {fun.__name__}() inside pure '
	f'{context.name}(). Is that intentional? If you want to nest pure '
	f'functions, use pure(..., nested=True) for the inner function.')
	# raise RuntimeError(
	# f'If you want to nest run() calls, use nested=True. ({context})')
	before = context
	try:
	name = fun.__name__
	if rng.shape == ():
	rng = jax.random.PRNGKey(rng)
	context = Context(state.copy(), rng, create, modify, ignore, [], name)
	CONTEXT[threading.get_ident()] = context
	out = fun(args, *kwargs)
	state = dict(context)
	return out, state
	finally:
	CONTEXT[threading.get_ident()] = before
	purified.pure = True
	return purified