Refactor interface for working with datasets.

[DavidIkov]

Current interface does not offer control, that some users may need for testing. For example you cant change inference size in a simple way.