Class LayerBlock

java.lang.Object
smile.deep.layer.LayerBlock
All Implemented Interfaces:
Function<Tensor,Tensor>, Layer
Direct Known Subclasses:
EfficientNet, FusedMBConv, LlamaModel, MBConv, MiniMaxModel, MiniMaxMoE, MiniMaxVisionTower, QwenModel, QwenMtp, QwenVisionTower, SequentialBlock, SqueezeExcitation
public abstract class LayerBlock extends Object implements Layer
A block is combinations of one or more layers. Blocks form the basis of more complex network designs. LayerBlock allows treating the whole container as a single layer, such that performing a transformation on the LayerBlock applies to each of the layers it contains (which are each a registered submodule of the block).

  • Field Details

    • module

      protected final MemorySegment module
      The neural network module (ST_Module).

    • device

      protected Device device
      The compute device.

    • dtype

      protected ScalarType dtype
      The data type.

  • Constructor Details

    • LayerBlock

      public LayerBlock()
      Constructor.

    • LayerBlock

      public LayerBlock(String name)
      Constructor.
      Parameters:
      name - the module name.

  • Method Details

    • module

      public MemorySegment module()
      Description copied from interface: Layer
      Returns the native ST_Module handle for this layer.
      Specified by:
      module in interface Layer
      Returns:
      the native ST_Module handle.

    • toString

      public String toString()
      Overrides:
      toString in class Object

    • add

      public LayerBlock add(String name, Layer layer)
      Adds a sub-layer.
      Parameters:
      name - the name of sub-layer.
      layer - the sub-layer.
      Returns:
      this object.

    • add

      public LayerBlock add(String name, MemorySegment layer)
      Adds a sub-module.
      Parameters:
      name - the name of sub-layer.
      layer - the native ST_Module handle of the sub-layer.
      Returns:
      this object.

    • isTraining

      public boolean isTraining()
      Returns true if the layer block is in training mode.
      Returns:
      true if the layer block is in training mode.

    • train

      public void train()
      Sets the layer block in the training mode.

    • eval

      public void eval()
      Sets the layer block in the evaluation/inference mode.

    • setRequiresGrad

      public LayerBlock setRequiresGrad(boolean required)
      Enables or disables requires_grad on every parameter (recursive). Call with false after loading an inference checkpoint so TP worker threads cannot accidentally build autograd graphs (NoGradGuard is thread-local and must also wrap each worker).
      Parameters:
      required - true to record gradients on parameters.
      Returns:
      this block.

    • device

      public Device device()
      Returns the compute device of module.
      Returns:
      the compute device of module.

    • dtype

      public ScalarType dtype()
      Returns the data type of module.
      Returns:
      the data type of module.

    • to

      public LayerBlock to(Device device)
      Description copied from interface: Layer
      Moves the layer block to a device.
      Specified by:
      to in interface Layer
      Parameters:
      device - the compute device.
      Returns:
      this layer.

    • to

      public LayerBlock to(Device device, ScalarType dtype)
      Description copied from interface: Layer
      Moves the layer block to a device.
      Specified by:
      to in interface Layer
      Parameters:
      device - the compute device.
      dtype - the data type.
      Returns:
      this layer.

    • load

      public void load(String path)
      Loads a checkpoint.
      Parameters:
      path - the checkpoint file path.

    • loadStateDict

      public void loadStateDict(Map<String,Tensor> stateDict, boolean strict)
      Loads parameters from a flat state dictionary whose keys match named_parameters() names (e.g. "layers.0.attention.wq.weight").
      Parameters:
      stateDict - map from fully-qualified parameter name to tensor.
      strict - when true, every module parameter must be present and every state-dict key must match a module parameter.

    • save

      public void save(String path)
      Serialize the layer block as a checkpoint.
      Parameters:
      path - the checkpoint file path.