def construct(
self,
query: ms.Tensor,
key: ms.Tensor,
value: ms.Tensor,
mask: Optional[ms.Tensor] = None) -> Tuple[ms.Tensor, ms.Tensor]:
batch_size = query.shape[0]
query = self.query(query)
key = self.key(key)
value = self.value(value)
# multi head
query = query.view(batch_size, -1, self.num_attention_heads,
self.dims_per_head).transpose(0, 2, 1, 3)
key = key.view(batch_size, -1, self.num_attention_heads,
self.dims_per_head).transpose(0, 2, 1, 3)
value = value.view(batch_size, -1, self.num_attention_heads,
self.dims_per_head).transpose(0, 2, 1, 3)
# self attention
context, attention = self.attention(query, key, value, attn_mask=mask)
# concat heads
context = context.transpose(0, 2, 1, 3).view(batch_size, -1,
self.hidden_size)
output = self.dense(context)
return output, attention
def MeanShift(x,
rgb_range,
rgb_mean=(0.4488, 0.4371, 0.4040),
rgb_std=(1.0, 1.0, 1.0),
sign=-1):
# super(MeanShift, self).__init__(3, 3, kernel_size=1)
std = Tensor(rgb_std)
conv2d = ops.Conv2D(out_channel=3, kernel_size=1)
biasadd = ops.BiasAdd()
weight = numpy.eye(3, 3).view((3, 3, 1, 1)) / std.view(3, 1, 1, 1)
bias = sign * rgb_range * Tensor(rgb_mean) / std
weight = weight.astype(numpy.float32)
bias = bias.astype(numpy.float32)
x = conv2d(x, weight)
x = biasadd(x, bias)
return x