Source code for ding.torch_utils.network.scatter_connection
import torch
import torch.nn as nn
from typing import Tuple, List
from ding.hpc_rl import hpc_wrapper
[docs]def shape_fn_scatter_connection(args, kwargs) -> List[int]:
"""
Overview:
Return the shape of scatter_connection for HPC.
Arguments:
- args (:obj:`Tuple`): The arguments passed to the scatter_connection function.
- kwargs (:obj:`Dict`): The keyword arguments passed to the scatter_connection function.
Returns:
- shape (:obj:`List[int]`): A list representing the shape of scatter_connection, \
in the form of [B, M, N, H, W, scatter_type].
"""
if len(args) <= 1:
tmp = list(kwargs['x'].shape)
else:
tmp = list(args[1].shape) # args[0] is __main__.ScatterConnection object
if len(args) <= 2:
tmp.extend(kwargs['spatial_size'])
else:
tmp.extend(args[2])
tmp.append(args[0].scatter_type)
return tmp
[docs]class ScatterConnection(nn.Module):
"""
Overview:
Scatter feature to its corresponding location. In AlphaStar, each entity is embedded into a tensor,
and these tensors are scattered into a feature map with map size.
Interfaces:
``__init__``, ``forward``, ``xy_forward``
"""
[docs] def __init__(self, scatter_type: str) -> None:
"""
Overview:
Initialize the ScatterConnection object.
Arguments:
- scatter_type (:obj:`str`): The scatter type, which decides the behavior when two entities have the \
same location. It can be either 'add' or 'cover'. If 'add', the first one will be added to the \
second one. If 'cover', the first one will be covered by the second one.
"""
super(ScatterConnection, self).__init__()
self.scatter_type = scatter_type
assert self.scatter_type in ['cover', 'add']
[docs] @hpc_wrapper(
shape_fn=shape_fn_scatter_connection,
namedtuple_data=False,
include_args=[0, 2],
include_kwargs=['x', 'location'],
is_cls_method=True
)
def forward(self, x: torch.Tensor, spatial_size: Tuple[int, int], location: torch.Tensor) -> torch.Tensor:
"""
Overview:
Scatter input tensor 'x' into a spatial feature map.
Arguments:
- x (:obj:`torch.Tensor`): The input tensor of shape `(B, M, N)`, where `B` is the batch size, `M` \
is the number of entities, and `N` is the dimension of entity attributes.
- spatial_size (:obj:`Tuple[int, int]`): The size `(H, W)` of the spatial feature map into which 'x' \
will be scattered, where `H` is the height and `W` is the width.
- location (:obj:`torch.Tensor`): The tensor of locations of shape `(B, M, 2)`. \
Each location should be (y, x).
Returns:
- output (:obj:`torch.Tensor`): The scattered feature map of shape `(B, N, H, W)`.
Note:
When there are some overlapping in locations, 'cover' mode will result in the loss of information.
'add' mode is used as a temporary substitute.
"""
device = x.device
B, M, N = x.shape
x = x.permute(0, 2, 1)
H, W = spatial_size
index = location[:, :, 1] + location[:, :, 0] * W
index = index.unsqueeze(dim=1).repeat(1, N, 1)
output = torch.zeros(size=(B, N, H, W), device=device).view(B, N, H * W)
if self.scatter_type == 'cover':
output.scatter_(dim=2, index=index, src=x)
elif self.scatter_type == 'add':
output.scatter_add_(dim=2, index=index, src=x)
output = output.view(B, N, H, W)
return output
[docs] def xy_forward(
self, x: torch.Tensor, spatial_size: Tuple[int, int], coord_x: torch.Tensor, coord_y
) -> torch.Tensor:
"""
Overview:
Scatter input tensor 'x' into a spatial feature map using separate x and y coordinates.
Arguments:
- x (:obj:`torch.Tensor`): The input tensor of shape `(B, M, N)`, where `B` is the batch size, `M` \
is the number of entities, and `N` is the dimension of entity attributes.
- spatial_size (:obj:`Tuple[int, int]`): The size `(H, W)` of the spatial feature map into which 'x' \
will be scattered, where `H` is the height and `W` is the width.
- coord_x (:obj:`torch.Tensor`): The x-coordinates tensor of shape `(B, M)`.
- coord_y (:obj:`torch.Tensor`): The y-coordinates tensor of shape `(B, M)`.
Returns:
- output (:obj:`torch.Tensor`): The scattered feature map of shape `(B, N, H, W)`.
Note:
When there are some overlapping in locations, 'cover' mode will result in the loss of information.
'add' mode is used as a temporary substitute.
"""
device = x.device
B, M, N = x.shape
x = x.permute(0, 2, 1)
H, W = spatial_size
index = (coord_x * W + coord_y).long()
index = index.unsqueeze(dim=1).repeat(1, N, 1)
output = torch.zeros(size=(B, N, H, W), device=device).view(B, N, H * W)
if self.scatter_type == 'cover':
output.scatter_(dim=2, index=index, src=x)
elif self.scatter_type == 'add':
output.scatter_add_(dim=2, index=index, src=x)
output = output.view(B, N, H, W)
return output