Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

GridMask implementation with BaseImageAugmentationLayer #159

Merged
merged 47 commits into from
Apr 1, 2022
Merged

GridMask implementation with BaseImageAugmentationLayer #159

merged 47 commits into from
Apr 1, 2022

Conversation

chjort
Copy link
Contributor

@chjort chjort commented Feb 28, 2022 •
edited
Loading

GridMask layer has been updated to use BaseImageAugmentationLayer, which closes #190.
As of now, auto-vectorization of GridMask is not functional using BaseImageAugmentationLayer.

Memory and speed optimizations to fill_utils.py are also included in this PR to improve GridMask performance.

chjort and others added 30 commits February 5, 2022 23:52
@chjort
temp driver
ba3413d
vectorize mask computation
a743946
vectorize mask computation
70587d9
vectorize mask computation
1fe49cd
comment
e18016c
random rotation and center cropping
8f03787
finish vectorized gridmask computation
6a0e9a6
refactor
27f2025
comment
92a019c
comment
15968f4
comments
2a24c94
@chjort
comments
fb443c7
@chjort
Merge branch 'master' into vectorize_gridmask
9edfdf3
@chjort
merge master into branch
abd8048
@chjort
initial vectorized layer
9599c15
@chjort
initial vectorized layer
ac8fa28
@chjort
debug memory usage
60759ce
@chjort
fix tests
fb6cfcd
@chjort
remove support for single image
79b8387
@chjort
set ratio random in demo
80c080a
@chjort
minimize memory by reducing number of simultaneous Logical And operat…
aa55af5 
…ions.
@chjort
individual ratio for each image when ratio="random"
a1d5f03
@chjort
add vectorized argument
3e8d0c4
@chjort
use float32 instead of int32
4a033ab
@chjort
use float32 instead of int32
de60f96
@chjort
remove vectorized arg
48df666
@chjort
minor refactor
a9f8700
@chjort
minor refactor
bbf8977
@chjort
support single image
54416ff
@chjort
refactor coordinates to mask
783b5c4
@chjort
Copy link
Contributor Author

chjort commented Feb 28, 2022 •
edited
Loading

Have you checked how the memory is going to increase with the input resolution other then the batch size? As I suppose in this version we need to store all the (dense) maskes in memory with the same input resolution for the whole batch.

Memory also scales with input resolution, but memory consumption does not come from storing the masks as is. It comes from the logical_and operation in this line which computes the masks.

I do not have any numbers for increasing image resolution. Let me know if I should conduct some experiments and collect some numbers. But memory will scale with O(batch_size * input_resolution)

@bhack
Copy link
Contributor

bhack commented Feb 28, 2022

Have you tried to jit_compile It?

@chjort
Copy link
Contributor Author

chjort commented Feb 28, 2022

Have you tried to jit_compile It?

Running the vectorized implementation with jit_compile=True on fill_utils.corners_to_mask gives minor speedup, but memory consumption remains the same as above:

Batch sizes Vectorized jit time [ms]
64 1725
128 1981
256 2375
512 3096

@bhack
Copy link
Contributor

bhack commented Feb 28, 2022

Probably the speed-up is a little bit more as we need to subtract the compile time.
I suppose that also in this version we cannot jit_compile the caller for the same reason as in #146 (comment) right?

@chjort chjort mentioned this pull request Mar 3, 2022
Closed
@LukeWood
Copy link
Contributor

LukeWood commented Mar 3, 2022

We are going to lean on the results from #160 to vectorize our KPLs. It looks like vectorized_map actually performs REALLY well in a tf.function, so we will be leaning on this. Let's see if we can rewrite this layer using the BaseImageAugmentationLayer.

I will be contributing an early sample for how to do this in the near future.

@LukeWood
Copy link
Contributor

@chjort the examples are ready. Check out RandomShear, AutoContrast, or RandomSharpness.

@chjort
Copy link
Contributor Author

chjort commented Mar 24, 2022 •
edited
Loading

Thanks @LukeWood, I will begin looking into this.

@chjort chjort marked this pull request as draft March 26, 2022 13:44
@chjort chjort mentioned this pull request Mar 30, 2022
Closed
@chjort chjort marked this pull request as ready for review March 30, 2022 18:05
qlzh727
qlzh727 requested changes Mar 30, 2022
View reviewed changes
Copy link
Member

@qlzh727 qlzh727 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR.

keras_cv/layers/preprocessing/grid_mask.py Outdated Show resolved Hide resolved
keras_cv/layers/preprocessing/grid_mask.py Outdated Show resolved Hide resolved
keras_cv/layers/preprocessing/grid_mask.py Outdated Show resolved Hide resolved
@chjort chjort changed the title Vectorized gridmask implementation GridMask implementation with BaseImageAugmentationLayer Mar 31, 2022
@chjort
Copy link
Contributor Author

chjort commented Apr 1, 2022

@qlzh727 Let me know if there are any other changes you would like me to do before this can be merged.

@innat
Copy link
Contributor

innat commented Apr 1, 2022

@chjort great work. Is it possible to showcase some demo output?

@chjort
Copy link
Contributor Author

chjort commented Apr 1, 2022 •
edited
Loading

@chjort great work. Is it possible to showcase some demo output?

Absolutely. Here are samples from GridMask(ratio="random", rotation_factor=0.5, fill_mode="gaussian_noise")
gridmask

And samples from GridMask(ratio=0.5, rotation_factor=0.5, fill_mode="constant", fill_value=0.0)
gridmask2

@LukeWood
Copy link
Contributor

LukeWood commented Apr 1, 2022

Thanks @chjort

@LukeWood LukeWood merged commit ca6eaa8 into keras-team:master Apr 1, 2022
ianstenbit pushed a commit to ianstenbit/keras-cv that referenced this pull request Aug 6, 2022
GridMask implementation with BaseImageAugmentationLayer (keras-team#159)
d514036 
* temp driver

* vectorize mask computation

* vectorize mask computation

* vectorize mask computation

* comment

* random rotation and center cropping

* finish vectorized gridmask computation

* refactor

* comment

* comment

* comments

* comments

* merge master into branch

* initial vectorized layer

* initial vectorized layer

* debug memory usage

* fix tests

* remove support for single image

* set ratio random in demo

* minimize memory by reducing number of simultaneous Logical And operations.

* individual ratio for each image when ratio="random"

* add vectorized argument

* use float32 instead of int32

* use float32 instead of int32

* remove vectorized arg

* minor refactor

* minor refactor

* support single image

* refactor coordinates to mask

* merge master

* optimize bounding box mask generation

* add tests

* minor refactor

* minor refactor

* formatting

* black

* merge with master

* WIP. GridMask to BaseImageAugmentationLayer

* WIP. GridMask to BaseImageAugmentationLayer

* merge master

* GridMask to BaseImageAugmentationLayer

* Apply changes from review

* fix test
adhadse pushed a commit to adhadse/keras-cv that referenced this pull request Sep 17, 2022
@chjort
GridMask implementation with BaseImageAugmentationLayer (keras-team#159)
f8b8472 
* temp driver

* vectorize mask computation

* vectorize mask computation

* vectorize mask computation

* comment

* random rotation and center cropping

* finish vectorized gridmask computation

* refactor

* comment

* comment

* comments

* comments

* merge master into branch

* initial vectorized layer

* initial vectorized layer

* debug memory usage

* fix tests

* remove support for single image

* set ratio random in demo

* minimize memory by reducing number of simultaneous Logical And operations.

* individual ratio for each image when ratio="random"

* add vectorized argument

* use float32 instead of int32

* use float32 instead of int32

* remove vectorized arg

* minor refactor

* minor refactor

* support single image

* refactor coordinates to mask

* merge master

* optimize bounding box mask generation

* add tests

* minor refactor

* minor refactor

* formatting

* black

* merge with master

* WIP. GridMask to BaseImageAugmentationLayer

* WIP. GridMask to BaseImageAugmentationLayer

* merge master

* GridMask to BaseImageAugmentationLayer

* Apply changes from review

* fix test
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@qlzh727 qlzh727 qlzh727 approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

Migrate GridMask to use BaseImageAugmentationLayer
5 participants
@chjort @bhack @LukeWood @innat @qlzh727