def __init__(self, input_shape, num_rotations, crop_size, preprocess): # pylint: disable=g-doc-args
"""Inits transport module with separate goal FCN.
Assumes the presence of a goal image, that cropping is done after the
query, that per-pixel loss is not used, and SE(2) grasping.
"""
self.num_rotations = num_rotations
self.crop_size = crop_size # crop size must be N*16 (e.g. 96)
self.preprocess = preprocess
self.lr = 1e-5
self.pad_size = int(self.crop_size / 2)
self.padding = np.zeros((3, 2), dtype=int)
self.padding[:2, :] = self.pad_size
input_shape = np.array(input_shape)
input_shape[0:2] += self.pad_size * 2
input_shape = tuple(input_shape)
# Output dimension (i.e., number of channels) of 3.
self.odim = output_dim = 3
# 3 fully convolutional ResNets. Third one is for the goal.
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix='s0_')
in1, out1 = ResNet43_8s(input_shape, output_dim, prefix='s1_')
in2, out2 = ResNet43_8s(input_shape, output_dim, prefix='s2_')
self.model = tf.keras.Model(inputs=[in0, in1, in2],
outputs=[out0, out1, out2])
self.optim = tf.keras.optimizers.Adam(learning_rate=self.lr)
self.metric = tf.keras.metrics.Mean(name='transport_loss')
def __init__(self, input_shape, d_action, use_mdn):
del use_mdn
self.batch_size = 4
self.input_shape = input_shape
self.d_action = d_action
channel_depth_dim = 16
in0, out0 = ResNet43_8s(self.input_shape,
channel_depth_dim,
prefix="s0_",
cutoff_early=False,
include_batchnorm=True)
# out0 = tf.nn.avg_pool(out0, ksize=(1,4,4,1), strides=(1,4,4,1),
# padding="SAME", data_format="NHWC")
out0 = compute_spatial_soft_argmax(out0, self.batch_size, 320, 160,
16) # shape (B, C, 2)
out0 = tf.keras.layers.Flatten()(out0)
out0 = tf.keras.layers.Dense(128,
kernel_initializer="normal",
bias_initializer="normal",
activation="relu")(out0)
out0 = tf.keras.layers.Dense(128,
kernel_initializer="normal",
bias_initializer="normal",
activation="relu")(out0)
out0 = tf.keras.layers.Dense(self.d_action,
kernel_initializer="normal",
bias_initializer="normal")(out0)
self.model = tf.keras.Model(inputs=[in0], outputs=[out0])
def __init__(self,
input_shape,
descriptor_dim,
num_rotations,
preprocess,
lite=False):
self.preprocess = preprocess
self.num_rotations = num_rotations
self.descriptor_dim = descriptor_dim
max_dim = np.max(input_shape[:2])
self.padding = np.zeros((3, 2), dtype=int)
pad = (max_dim - np.array(input_shape[:2])) / 2
self.padding[:2] = pad.reshape(2, 1)
input_shape = np.array(input_shape)
input_shape += np.sum(self.padding, axis=1)
input_shape = tuple(input_shape)
# Initialize fully convolutional Residual Network with 43 layers and
# 8-stride (3 2x2 max pools and 3 2x bilinear upsampling)
if lite:
d_in, d_out = ResNet36_4s(input_shape, self.descriptor_dim)
else:
d_in, d_out = ResNet43_8s(input_shape, self.descriptor_dim)
self.model = tf.keras.models.Model(inputs=[d_in], outputs=[d_out])
self.optim = tf.keras.optimizers.Adam(learning_rate=1e-5)
self.metric = tf.keras.metrics.Mean(name='attention_loss')
def __init__(self, in_shape, n_rotations, preprocess, lite=False):
self.n_rotations = n_rotations
self.preprocess = preprocess
max_dim = np.max(in_shape[:2])
self.padding = np.zeros((3, 2), dtype=int)
pad = (max_dim - np.array(in_shape[:2])) / 2
self.padding[:2] = pad.reshape(2, 1)
in_shape = np.array(in_shape)
in_shape += np.sum(self.padding, axis=1)
in_shape = tuple(in_shape)
# Initialize fully convolutional Residual Network with 43 layers and
# 8-stride (3 2x2 max pools and 3 2x bilinear upsampling)
if lite:
d_in, d_out = ResNet36_4s(in_shape, 1)
else:
d_in, d_out = ResNet43_8s(in_shape, 1)
self.model = tf.keras.models.Model(inputs=[d_in], outputs=[d_out])
self.optim = tf.keras.optimizers.Adam(learning_rate=1e-4)
# self.optim = tf.keras.optimizers.SGD(learning_rate=1e-3, momentum=0.9, nesterov=True, name='SGD')
self.metric = tf.keras.metrics.Mean(name='loss_attention')
def __init__(self, in_shape, n_rotations, crop_size, preprocess):
"""Transport module for placing.
Args:
in_shape: shape of input image.
n_rotations: number of rotations of convolving kernel.
crop_size: crop size around pick argmax used as convolving kernel.
preprocess: function to preprocess input images.
"""
self.iters = 0
self.n_rotations = n_rotations
self.crop_size = crop_size # crop size must be N*16 (e.g. 96)
self.preprocess = preprocess
self.pad_size = int(self.crop_size / 2)
self.padding = np.zeros((3, 2), dtype=int)
self.padding[:2, :] = self.pad_size
in_shape = np.array(in_shape)
in_shape[0:2] += self.pad_size * 2
in_shape = tuple(in_shape)
# Crop before network (default for Transporters in CoRL submission).
kernel_shape = (self.crop_size, self.crop_size, in_shape[2])
if not hasattr(self, 'output_dim'):
self.output_dim = 3
if not hasattr(self, 'kernel_dim'):
self.kernel_dim = 3
# 2 fully convolutional ResNets with 57 layers and 16-stride
in0, out0 = ResNet43_8s(in_shape, self.output_dim, prefix='s0_')
# in1, out1 = ResNet43_8s(in_shape, self.kernel_dim, prefix='s1_')
in1, out1 = ResNet43_8s(kernel_shape, self.kernel_dim, prefix='s1_')
self.model = tf.keras.Model(inputs=[in0, in1], outputs=[out0, out1])
self.optim = tf.keras.optimizers.Adam(learning_rate=1e-4)
# self.optim = tf.keras.optimizers.SGD(learning_rate=1e-4, momentum=0.9, nesterov=True, name='SGD')
self.metric = tf.keras.metrics.Mean(name='loss_transport')
