neuroailab / tfutils

Utilities for working with tensorflow
MIT License
25 stars 8 forks source link

Uses a single parameter to toggle between using tpu or not #133

Closed jvrsgsty closed 5 years ago

jvrsgsty commented 5 years ago

I found it weird that to use this unified data provider on tpu you have to specify you are doing things on tpu twice: once with the on_tpu arg in the constructor, and once when you choose which function to pass in the data params dictionary.

I also thought it would be best if we keep things consistent between gpu and tpu and I removed is_train from the constructor args, and pass it via the kwargs in the dataset func.

With this, usage would go from

'data_params': {
            'func': ImageNet(image_dir=FLAGS.data_dir,
                             prep_type='resnet',
                             on_tpu=True,
                             is_train=True,
                             resize=FLAGS.image_size).dataset_func_tpu,
            'batch_size': FLAGS.train_batch_size
        }

to this, which would only differ from a GPU call in the on_tpu argument.

'data_params': {
            'func': ImageNet(image_dir=FLAGS.data_dir,
                             prep_type='resnet',
                             on_tpu=True,
                             resize=FLAGS.image_size).dataset_func,
            'is_train': True,
            'batch_size': FLAGS.train_batch_size
        }
jvrsgsty commented 5 years ago

Ugh... Basically had to roll back to almost as this was before. The reason we are doing this this way, is that the TPUEstimator API wants the dataset func to strictly adhere to this signature: it must have a positional argument named params.

So what I propose instead, is inferring self.on_tpu depending on which dataset_func is called.

Sadly, we do have to adhere to two different call signatures, for now.

anayebi commented 5 years ago

This makes sense to me, merging now.