wehub-resource-sync
86 lines
2.9 KiB
Python
86 lines
2.9 KiB
Python
from ray.rllib.utils.framework import try_import_torch
|
|
|
|
torch, nn = try_import_torch()
|
|
|
|
|
|
class Stride2D(nn.Module):
|
|
"""A striding layer for doing torch Conv2DTranspose operations.
|
|
|
|
Using this layer before the 0-padding (on a 3D input "image") and before
|
|
the actual ConvTranspose2d allows for a padding="same" behavior that matches
|
|
100% that of a `tf.keras.layers.Conv2DTranspose` layer.
|
|
|
|
Examples:
|
|
Input image (4x4):
|
|
A B C D
|
|
E F G H
|
|
I J K L
|
|
M N O P
|
|
|
|
Stride with stride=2 -> output image=(7x7)
|
|
A 0 B 0 C 0 D
|
|
0 0 0 0 0 0 0
|
|
E 0 F 0 G 0 H
|
|
0 0 0 0 0 0 0
|
|
I 0 J 0 K 0 L
|
|
0 0 0 0 0 0 0
|
|
M 0 N 0 O 0 P
|
|
"""
|
|
|
|
def __init__(self, width, height, stride_w, stride_h):
|
|
"""Initializes a Stride2D instance.
|
|
|
|
Args:
|
|
width: The width of the 3D input "image".
|
|
height: The height of the 3D input "image".
|
|
stride_w: The stride in width direction, with which to stride the incoming
|
|
image.
|
|
stride_h: The stride in height direction, with which to stride the incoming
|
|
image.
|
|
"""
|
|
super().__init__()
|
|
|
|
self.width = width
|
|
self.height = height
|
|
self.stride_w = stride_w
|
|
self.stride_h = stride_h
|
|
|
|
self.register_buffer(
|
|
"zeros",
|
|
torch.zeros(
|
|
size=(
|
|
self.width * self.stride_w - (self.stride_w - 1),
|
|
self.height * self.stride_h - (self.stride_h - 1),
|
|
),
|
|
dtype=torch.float32,
|
|
),
|
|
)
|
|
|
|
self.out_width, self.out_height = self.zeros.shape[0], self.zeros.shape[1]
|
|
# Squeeze in batch and channel dims.
|
|
self.zeros = self.zeros.unsqueeze(0).unsqueeze(0)
|
|
|
|
where_template = torch.zeros(
|
|
(self.stride_w, self.stride_h), dtype=torch.float32
|
|
)
|
|
# Set upper/left corner to 1.0.
|
|
where_template[0][0] = 1.0
|
|
# then tile across the entire (strided) image size.
|
|
where_template = where_template.repeat((self.height, self.width))[
|
|
: -(self.stride_w - 1), : -(self.stride_h - 1)
|
|
]
|
|
# Squeeze in batch and channel dims and convert to bool.
|
|
where_template = where_template.unsqueeze(0).unsqueeze(0).bool()
|
|
self.register_buffer("where_template", where_template)
|
|
|
|
def forward(self, x):
|
|
# Repeat incoming image stride(w/h) times to match the strided output template.
|
|
repeated_x = (
|
|
x.repeat_interleave(self.stride_w, dim=-2).repeat_interleave(
|
|
self.stride_h, dim=-1
|
|
)
|
|
)[:, :, : -(self.stride_w - 1), : -(self.stride_h - 1)]
|
|
# Where `self.where_template` == 1.0 -> Use image pixel, otherwise use
|
|
# zero filler value.
|
|
return torch.where(self.where_template, repeated_x, self.zeros)
|