added a jax example #11
Conversation
tanayvarshney
force-pushed
the
tvarshney_jax
branch
from
0d26c62
to
f2ed614
Compare
| { | ||
| name: "fc_out" | ||
| data_type: TYPE_FP32 | ||
| dims: [-1, -1] |
There was a problem hiding this comment.
This doesn't seem right. Shouldn't it be dims: [1000]?
There was a problem hiding this comment.
It [-1,1000] if I remember it correctly. I can re-run and check
Quick_Deploy/JAX/README.md
Outdated
| ``` | ||
| class TritonPythonModel: | ||
| def initialize(self, args): | ||
| ... |
There was a problem hiding this comment.
Can you add some explanation of what you are putting in the config.pbtxt file? Mainly the input/output names and how they have to correspond to what's in the execute method
There was a problem hiding this comment.
We have this info in other places, I am trying to avoid replication of information, but I see why it is needed. We should discuss more about this. We are bound to run into this issue later down the line
| inp = pb_utils.get_input_tensor_by_name(request, "image") | ||
| input_image = inp.as_numpy() | ||
|
|
||
| params = self.resnet18.init(self.key, input_image) |
There was a problem hiding this comment.
I'm surprised by this init call here. I'm not overly familiar with how Flax or the Flaxmodels library structure things, but why do we call init here instead of in the initialize method. What does it actually do? Is it about needing to know the shape of the input image?
There was a problem hiding this comment.
Ah, good catch, I missed this.
|
Pinging @tanayvarshney as a reminder, in case this is still in progress. |
No description provided.