mindflow.cell.AttentionBlock
- class mindflow.cell.AttentionBlock(in_channels, num_heads, drop_mode='dropout', dropout_rate=0.0, compute_dtype=mstype.float32)[source]
AttentionBlock comprises an MultiHeadAttention and an MLP layer.
- Parameters
in_channels (int) – The input channels.
num_heads (int) – The number of attention heads.
drop_mode (str) – Dropout method. Default:
dropout
. Supportdropout
ordroppath
.dropout_rate (float) – The drop rate of dropout layer, greater than 0 and less equal than 1. Default:
0.0
.compute_dtype (mindspore.dtype) – Compute dtype. Default:
mstype.float32
, indicatesmindspore.float32
.
- Inputs:
x (Tensor) - Tensor with shape
.mask (Tensor) - Tensor with shape
or or .
- Outputs:
output (Tensor) - Tensor with shape
.
- Supported Platforms:
Ascend
CPU
Examples
>>> from mindspore import ops >>> from mindflow.cell import AttentionBlock >>> model = AttentionBlock(in_channels=256, num_heads=4) >>> x = ops.rand((4, 100, 256)) >>> output = model(x) >>> print(output.shape) (4, 100, 256)