-
Notifications
You must be signed in to change notification settings - Fork 330
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.
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
Add GPU testing for Jax+Torch #1935
Changes from all commits
6087fbd
6a1eae3
f885ddb
fdf201c
9066e3f
16ef366
8fb05ea
f66a72c
6a55cfb
c41ab45
48d3039
dac2df3
c8886b9
039d822
f4e2962
ed1d6e0
32fd07a
c76e352
1a4a0e1
b5ee69a
495e6ed
24aaa65
a746374
f5c74e0
6491da2
dd6730f
91760f9
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,4 +1,4 @@ | ||
# keras-cv-image:deps has all deps of KerasCV for testing. | ||
FROM us-west1-docker.pkg.dev/keras-team-test/keras-cv-test/keras-cv-image:deps | ||
ARG IMAGE_NAME | ||
FROM $IMAGE_NAME | ||
COPY . /kerascv | ||
WORKDIR /kerascv |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -4,10 +4,11 @@ local gpus = import 'templates/gpus.libsonnet'; | |
local image = std.extVar('image'); | ||
local tagName = std.extVar('tag_name'); | ||
local gcsBucket = std.extVar('gcs_bucket'); | ||
local backend = std.extVar('backend'); | ||
|
||
local unittest = base.BaseTest { | ||
// Configure job name. | ||
frameworkPrefix: "tf", | ||
frameworkPrefix: backend, | ||
modelName: "keras-cv", | ||
mode: "unit-tests", | ||
timeout: 3600, # 1 hour, in seconds | ||
|
@@ -21,20 +22,28 @@ local unittest = base.BaseTest { | |
entrypoint: [ | ||
'bash', | ||
'-c', | ||
||| | ||
# Build custom ops from source | ||
python build_deps/configure.py | ||
bazel-5.4.0 build keras_cv/custom_ops:all --verbose_failures | ||
cp bazel-bin/keras_cv/custom_ops/*.so keras_cv/custom_ops/ | ||
export TEST_CUSTOM_OPS=true | ||
std.format( | ||
||| | ||
export KERAS_BACKEND=%s | ||
export JAX_ENABLE_X64=true | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why do you end up needed int64/float64? Just curious really. In KerasNLP we have just been trying to go int32 everywhere by default. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think we probably should just do int32 everywhere. During NMS porting I had some things that were real sticklers about int64 in TF and I never found a way to work around There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could we potentially just change the dtype for those tests? Asking since we're always flirting with OOM issues and int64 can't help There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Potentially, yes. I did some looking and it seems like the only place we're using int64 internally with KerasCore is in the YOLOV8 label encoder -- I guess I already got rid of them from NMS. I'll check with Tirth if he's planning on getting rid of those. |
||
|
||
# Run whatever is in `command` here. | ||
${@:0} | ||
||| | ||
# Run whatever is in `command` here. | ||
${@:0} | ||
|||, | ||
backend | ||
) | ||
], | ||
command: [ | ||
'pytest --run_large --durations 0', | ||
'keras_cv', | ||
'keras_cv/bounding_box', | ||
'keras_cv/callbacks', | ||
'keras_cv/losses', | ||
'keras_cv/layers/object_detection', | ||
'keras_cv/layers/preprocessing', | ||
'keras_cv/models/backbones', | ||
'keras_cv/models/classification', | ||
'keras_cv/models/object_detection/retinanet', | ||
'keras_cv/models/object_detection/yolo_v8', | ||
], | ||
}; | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I assume
--run_large
was a mistake here since this is CI?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not a mistake, I was just doing this during development of the rebase so that we covered all the tests. But yes we shouldn't include it anymore