Skip to content

Allow taking ownership of convolution weights#32

Merged
mtavenrath merged 1 commit intorustnn:mainfrom
theHamsta:conv-weight-copied
Apr 1, 2026
Merged

Allow taking ownership of convolution weights#32
mtavenrath merged 1 commit intorustnn:mainfrom
theHamsta:conv-weight-copied

Conversation

@theHamsta
Copy link
Copy Markdown
Contributor

@theHamsta theHamsta commented Apr 1, 2026

feat: allow to take ownership of ConvWeights

It's difficult for users of trtx to create temporary weights for convolutions
if they decide to perform the swizzling themselves (they could also swizzle using TensorRT layers).

The ownership model expects that the user can prove that weights remain
valid throughout the whole lifetime of NetworkDefinition. This is
typically done by non-exclusively borrowing from weights from somewhere.

If they want to create new weights they typically need a mutual borrow and
unsafe code as a mutual borrow could also delete weights prematurely.

Since we already have a storage for owned small weights already we can store such
weights on behave of the user and ensure the weights stay alive during
the lifetime of NetworkDefinition.

@theHamsta theHamsta changed the title Conv weight copied Allow taking ownership of convolution weights Apr 1, 2026
@theHamsta theHamsta force-pushed the conv-weight-copied branch 5 times, most recently from 15e2ceb to 34c86d7 Compare April 1, 2026 14:55
It's difficult for users of trtx to create temporary weights for convolutions
if they decide to perform the swizzling themselves (they could also swizzle using TensorRT layers).

The ownership model expects that the user can prove that weights remain
valid throughout the whole lifetime of `NetworkDefinition`. This is
typically done by non-exclusively borrowing from weights from somewhere.

If they want to create new weights they typically need a mutual borrow and
unsafe code as a mutual borrow could also delete weights prematurely.

Since we already have a storage for owned small weights already we can store such
weights on behave of the user and ensure the weights stay alive during
the lifetime of `NetworkDefinition`.
@theHamsta theHamsta force-pushed the conv-weight-copied branch from 34c86d7 to dec817b Compare April 1, 2026 15:07
@theHamsta theHamsta marked this pull request as ready for review April 1, 2026 16:49
Comment thread trtx/src/network.rs
/// Same as [`Self::add_convolution`] but allows to set weights later (e.g. dynamic kernel and bias)
///
/// See [`trtx_sys::nvinfer1::INetworkDefinition::addConvolutionNd`].
pub fn add_convolution_deferrred_weights(
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

we could not add this method and then also make kernel weights optional. but maybe an explicit function for that gives a hint that this is possible at all

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

It looks like a deficit in the original TRT-RTX interface. Your solution is a good workaround for now.

@theHamsta theHamsta requested a review from mtavenrath April 1, 2026 16:57
@mtavenrath mtavenrath merged commit 9ccf4ea into rustnn:main Apr 1, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants