diff options
Diffstat (limited to 'model/stream.py')
| -rw-r--r-- | model/stream.py | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/model/stream.py b/model/stream.py index 61ff1c3..d69b962 100644 --- a/model/stream.py +++ b/model/stream.py @@ -1,4 +1,4 @@ -from fuel.transformers import Batch, Padding, Mapping, SortMapping, Unpack, MultiProcessing +from fuel.transformers import Batch, Padding, Mapping, SortMapping, Unpack, MultiProcessing, Filter from fuel.streams import DataStream from fuel.schemes import ConstantScheme, ShuffledExampleScheme @@ -31,6 +31,12 @@ class StreamRec(object): elif not data.tvt: stream = transformers.add_destination(stream) + if hasattr(self.config, 'train_max_len'): + idx = stream.sources.index('latitude') + def max_len_filter(x): + return len(x[idx]) <= self.config.train_max_len + stream = Filter(stream, max_len_filter) + stream = transformers.TaxiExcludeEmptyTrips(stream) stream = transformers.taxi_add_datetime(stream) stream = transformers.Select(stream, tuple(v for v in req_vars if not v.endswith('_mask'))) |