Skip to content

Add NASNet-a model - #1688

Open
easyeasydev wants to merge 4 commits into
flexflow:masterfrom
easyeasydev:add-nasnet-model
Open

easyeasydev wants to merge 4 commits into
flexflow:masterfrom
easyeasydev:add-nasnet-model

Conversation

@easyeasydev

@easyeasydev easyeasydev commented Sep 29, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Add NASNet-A Large to models.

The implementation includes:

  • NASNet stem, normal, reduction, and factorized-reduction cells
  • Config generated through dtgen
  • TensorFlow-style SAME padding for pooling and depthwise convolutions
  • Correct shifted path in factorized reduction
  • Model configuration and graph-structure tests
  • nasnet_a_large support in export-model-arch

PCG and task-spec changes

NASNet requires asymmetric padding and cropping, but the current computation graph does not expose dedicated padding or slicing operators. The model therefore composes these operations using existing split, concat, and scalar operators.

This exposed two existing split infrastructure issues:

  • ComputationGraphBuilder::split validated the requested output count against variadic input slots instead of output slots.
  • The split task implementation retrieved a single OUTPUT tensor even though split produces OUTPUT_00, OUTPUT_01, etc. This resulted in incorrect output pointers and block sizes for multi-output execution.

This PR corrects the PCG validation and updates the split forward/backward task implementations to collect every variadic output independently. It also adds coverage for rejecting unsupported output counts.


This change is Reviewable

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.

1 participant