Two Branches Support after Final Convolution Layer

[wolgwang1729]

Pull Request: Two Branches Support after Final Convolution Layer

Summary

This pull request adds support for dual branches in the AlexNet visualization after the final convolutional layer. Users can now dynamically add, remove, and visualize two independent branches ("Branch 1" and "Branch 2"). This is useful for architectures with multi-task heads, auxiliary classifiers, or Siamese/parallel branches.
Some papers that use this architecture include:

• Pose Estimation for Non-Cooperative Rendezvous Using Neural Networks
• Multi-Task Learning with Deep Neural Networks: A Survey
• Multi-task Learning For Joint Action and Gesture Recognition

Changes

HTML (AlexNet.html)

• Branch UI Additions:
  • Added two sections within #architecture2:
    • .branch-1 for Branch 1 layers
    • .branch-2 for Branch 2 layers
  • Each branch has its own input group and add/remove buttons.
  • Clear labels for each branch.

JavaScript (AlexNet.js)

• Data Structure Update:
  • Changed architecture2 from an array to an object: { branch1: [], branch2: [] }.
• Layer Parsing:
  • On redraw, each branch's layers are parsed independently from their respective inputs.
• 3D Visualization Logic:
  • Each branch's layers are rendered as separate stacks, offset in the X-axis.
  • Arrows connect the last convolution layer to each branch.
  • Handles cases where only one or both branches are present.
• UI Behavior:
  • Add/remove buttons work independently for each branch.
  • Visualization updates live as the user edits layers.

UI Behavior

Dynamic Interface Elements

• Add/Remove Layer:
  • Each branch has its own "+" (add) and "−" (remove) buttons.
  • Layers can be added/removed dynamically in either branch.
• Live Updates:
  • Visualization updates in real-time as the user edits layer configurations.

User Workflow

1. Define Main Architecture:
   Add convolutional/dense layers to the main architecture as before.
2. Add Branches:
   Add layers to either or both branches after the final convolution layer.
3. Visualize:
   Visualization shows the main trunk and both branches, clearly separated and connected.

Images

Dual Branch Neural Network:

Sidebar:

Notes:

• I ensure that the new branches do not interfere with the main architecture. So if branch 2 is empty, it will not affect the visualization of branch 1 and it is exactly like before.
• I don't know if it should be added to the main branch or not. I needed this feature for my own project, so I thought it would be useful to others as well.

Further Improvements

• I tried generalising the code to support more than two branches, but it became too complex and require significant logic for angle and distance calculations. Moreover more than 5 branches would be too complex to visualize clearly in my opinion. So if someone wants to add more branches, they can do so by duplicating the existing branch code and updating the logic accordingly.
• Better implementation. Currenly, the code is a bit repetitive for each branch. A more modular approach could be implemented to reduce redundancy and a better UI of the sidebar so that Branch 2 column is not visible on the outset.
• Allowing users to specify branch names.
• Convolutional layers in each branch.

[alexlenail]

Thanks for this other PR too!

• Are the branches always equal size? Could there be just a checkbox for "add a second fully connected branch" ? Or does the flexibility really matter here?
• In general we should avoid including niche models on the site (but let people customize like you have). Do you think this model architecture is too niche for the default site, and should keep it as a branch? Or do you think it's common enough others will use it and we should merge?

[wolgwang1729]

Thanks for the suggestions:

• No, the branches are not always of equal size, so I think the flexibility is important. I created the checkbox to add a second branch so that the initial UI shows only one branch by default, but users can add a second branch if they want to.
• I have found this type of model many times.
  • Nobody uses AlexNet in isolation these days. Modern workflows frequently reuse an AlexNet-style backbone (or any CNN backbone) to extract features and attach two different heads for example, one for classification and one for a regression or similarity task or they can attack some other layer like transformer.
  • Giving this second branch option reduces their work and efforts as they can create a dummy representation of that layer, download the SVG and just label the,m which will save their time.

As I said, I had modified this, so I just wanted to share if there is someone who really wants to do something similar and check this PR. I’m fine if you decide to keep it as an optional extension rather than merging into the default site. The goal was simply to contribute a feature that already saved me time; if it can save others a few hours of diagram hacks, that’s a win.

[alexlenail]

Would you be willing to take a pass at resolving the conflicts that arose from merging #68?

[wolgwang1729]

I have resolved the conflicts.