def __init__(self,
name,
dim,
radius_init,
max_classes=20,
fix_unknown=False,
unknown_id=None,
similarity="euclidean",
static_beta_gamma=True,
unknown_logits="radii",
radius_init_write=None,
use_ssl_beta_gamma_write=True,
temp_init=10.0,
dtype=tf.float32):
assert unknown_logits == 'radii'
super(SemiSupervisedMinDistGRUProtoMemory,
self).__init__(name,
dim,
radius_init,
max_classes=max_classes,
fix_unknown=fix_unknown,
unknown_id=unknown_id,
similarity=similarity,
unknown_logits=unknown_logits,
temp_init=temp_init,
dtype=dtype)
self._radius_init = radius_init
log.info('Radius init {}'.format(radius_init))
if radius_init_write is not None:
self._radius_init_write = radius_init_write
log.info('Radius init write {}'.format(radius_init_write))
else:
self._radius_init_write = radius_init
self._use_ssl_beta_gamma_write = use_ssl_beta_gamma_write
with variable_scope(name):
self._storage = GRU1DMod("storage",
dim,
dim,
layernorm=False,
dtype=dtype)
if static_beta_gamma:
with variable_scope(name):
self._beta = self._get_variable(
"beta", self._get_constant_init([], radius_init))
self._gamma = self._get_variable(
"gamma", self._get_constant_init([], 1.0))
self._beta2 = self._get_variable(
"beta2",
self._get_constant_init([], self._radius_init_write))
self._gamma2 = self._get_variable(
"gamma2", self._get_constant_init([], 1.0))
def __init__(self,
name,
in_filter,
out_filter,
stride,
data_format="NCHW",
leaky_relu=0.0,
dtype=tf.float32):
super(ResidualSnailModule, self).__init__()
self._data_format = data_format
self._stride = stride
self._leaky_relu = leaky_relu
with variable_scope(name):
self._conv1 = Conv2D(
"conv1", 3, in_filter, out_filter, self._stride_arr(1), dtype=dtype)
self._bn1 = BatchNorm(
"bn1", out_filter, dtype=dtype, data_format=data_format)
self._conv2 = Conv2D(
"conv2", 3, out_filter, out_filter, self._stride_arr(1), dtype=dtype)
self._bn2 = BatchNorm(
"bn2", out_filter, dtype=dtype, data_format=data_format)
self._conv3 = Conv2D(
"conv3", 1, out_filter, out_filter, self._stride_arr(1), dtype=dtype)
self._bn3 = BatchNorm(
"bn3", out_filter, dtype=dtype, data_format=data_format)
self._ds = stride[2] > 1 or in_filter != out_filter
if self._ds:
self._projconv = Conv2D(
"projconv",
1,
in_filter,
out_filter,
self._stride_arr(1),
dtype=dtype)
def __init__(self,
name,
dim,
radius_init,
max_classes=20,
fix_unknown=False,
unknown_id=None,
similarity="euclidean",
static_beta_gamma=True,
unknown_logits="radii",
temp_init=10.0,
dtype=tf.float32,
**kwargs):
super(MinDistProtoMemory, self).__init__(name,
dim,
max_classes=max_classes,
fix_unknown=fix_unknown,
unknown_id=unknown_id,
similarity=similarity,
temp_init=temp_init,
dtype=dtype)
self._radius_init = radius_init
self._unknown_logits = unknown_logits
if static_beta_gamma:
with variable_scope(name):
self._beta = self._get_variable(
"beta", self._get_constant_init([], radius_init))
self._gamma = self._get_variable(
"gamma", self._get_constant_init([], 1.0))
def __init__(self,
name,
in_filter,
out_filter,
stride=2,
add_relu=True,
data_format="NCHW",
pool_padding="SAME",
dtype=tf.float32,
wdict=None):
super(ConvModule, self).__init__()
self._data_format = data_format
with variable_scope(name):
self._conv = Conv2D("conv",
3,
in_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype,
wdict=wdict)
self._bn = BatchNorm("bn",
out_filter,
data_format=data_format,
dtype=dtype,
wdict=wdict)
self._stride = stride
self._add_relu = add_relu
self._pool_padding = pool_padding
def __init__(self,
name,
dim,
max_items=80,
max_classes=20,
unknown_id=None,
log_sigma_init=0.0,
log_lambda_init=0.0,
radius_init=10.0,
similarity="euclidean",
dtype=tf.float32):
super(ExampleMemory, self).__init__(dtype=dtype)
sigma_init = self._get_constant_init([], log_sigma_init)
lbd_init = self._get_constant_init([], log_lambda_init)
self._similarity = similarity
self._dim = dim
with variable_scope(name):
self._log_sigma_u = self._get_variable("sigma_u", sigma_init)
self._log_sigma_l = self._get_variable("sigma_l", sigma_init)
self._beta = self._get_variable(
"beta", self._get_constant_init([], radius_init))
self._gamma = self._get_variable("gamma",
self._get_constant_init([], 1.0))
assert unknown_id is not None, 'Need to provide unknown ID'
self._unknown_id = unknown_id
self._max_classes = max_classes
self._max_items = max_items
def __init__(self, name, in_dim, label_dim, config, dtype=tf.float32):
super(DNCWriteHeadFeed2, self).__init__(name,
in_dim,
config,
dtype=dtype)
self._label_dim = label_dim
use_mlp = False
with variable_scope(name):
if use_mlp:
self._write_query_mlp = ResMLP('write_query_mlp', [
in_dim, self._nwrite * self._memory_dim, self._nwrite *
self._memory_dim, self._nwrite * self._memory_dim
],
dtype=dtype)
self._write_content_mlp = ResMLP('write_content_mlp', [
in_dim, self._nwrite * self._memory_dim, self._nwrite *
self._memory_dim, self._nwrite * self._memory_dim
],
dtype=dtype)
else:
self._write_query_mlp = Linear('write_query_mlp', in_dim,
self._nwrite * self._memory_dim)
self._write_content_mlp = Linear(
'write_content_mlp', in_dim,
self._nwrite * self._memory_dim)
def __init__(self,
name,
filter_size,
in_filter,
out_filter,
num_groups=32,
max_pool=True,
data_format="NCHW",
dtype=tf.float32):
super(InitConvModule, self).__init__()
self._data_format = data_format
self._max_pool = max_pool
with variable_scope(name):
self._conv = Conv2D("conv",
filter_size,
in_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._gn = GroupNorm("gn",
out_filter,
num_groups,
data_format=data_format,
dtype=dtype)
self._data_format = data_format
def build_net(self, wdict=None):
with variable_scope(self._name):
K = self.config.num_classes
if self.config.inner_loop_loss == "mix":
K += 1
if self.config.cosine_classifier:
# mlp = CosineLinear(
# "mlp",
# # list(self.config.num_filters) + [K],
# self.config.num_filters[0],
# K,
# temp=10.0,
# wdict=wdict)
mlp = CosineLastMLP(
"mlp",
list(self.config.num_filters) + [K],
# self.config.num_filters[0], K,
temp=10.0,
wdict=wdict)
else:
mlp = MLP("mlp",
list(self.config.num_filters) + [K],
wdict=wdict,
add_bias=self.config.classifier_bias)
return mlp
def __init__(self, name, nin, nout, dtype=tf.float32):
super(LSTM1DMod, self).__init__(dtype=dtype)
self._nin = nin
self._nout = nout
with variable_scope(name):
self._gates = Linear("gates_linear", nin + nout, nout + 2)
def __init__(self,
name,
in_filter,
out_filter,
num_groups,
stride=2,
data_format="NCHW",
dtype=tf.float32,
wdict=None):
super(ConvGNModule, self).__init__()
self._data_format = data_format
with variable_scope(name):
self._conv = Conv2D("conv",
3,
in_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype,
wdict=wdict)
self._gn = GroupNorm("gn",
out_filter,
num_groups,
data_format=data_format,
dtype=dtype,
wdict=wdict)
self._stride = stride
def __init__(self,
name,
in_filter,
out_filter,
stride,
data_format="NCHW",
dtype=tf.float32,
add_relu=True):
super(ResidualModule, self).__init__()
self._data_format = data_format
self._name = name
self._stride = stride
with variable_scope(name):
self._conv1 = Conv2D("conv1",
3,
in_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._bn1 = BatchNorm("bn1",
out_filter,
data_format=data_format,
dtype=dtype)
self._conv2 = Conv2D("conv2",
3,
out_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._bn2 = BatchNorm("bn2",
out_filter,
data_format=data_format,
dtype=dtype)
self._conv3 = Conv2D("conv3",
3,
out_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._bn3 = BatchNorm("bn3",
out_filter,
data_format=data_format,
dtype=dtype)
self._projconv = Conv2D(self._prefix(name, "projconv"),
1,
in_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._projbn = BatchNorm(self._prefix(name, "projbn"),
out_filter,
data_format=data_format,
dtype=dtype)
self._data_format = data_format
self._add_relu = add_relu
def __init__(self, name, num_channels, eps=1e-6, dtype=tf.float32):
super(LayerNorm, self).__init__(dtype=dtype)
shape = [num_channels]
with variable_scope(name):
binit = self._get_constant_init(shape, 0.0)
ginit = self._get_constant_init(shape, 1.0)
self._beta = self._get_variable("beta", binit)
self._gamma = self._get_variable("gamma", ginit)
self._eps = eps
def __init__(self,
name,
in_filter,
out_filter,
stride,
num_groups=32,
data_format="NCHW",
dtype=tf.float32,
add_relu=True):
super(BottleneckResidualModule, self).__init__()
self._data_format = data_format
self._stride = stride
with variable_scope(name):
self._conv1 = Conv2D("conv1",
1,
in_filter,
out_filter // 4,
stride,
dtype=dtype)
self._gn1 = GroupNorm("gn1",
out_filter // 4,
num_groups,
dtype=dtype)
self._conv2 = Conv2D("conv2",
3,
out_filter // 4,
out_filter // 4,
self._stride_arr(1),
dtype=dtype)
self._gn2 = GroupNorm("gn2",
out_filter // 4,
num_groups,
dtype=dtype)
self._conv3 = Conv2D("conv3",
1,
out_filter // 4,
out_filter,
self._stride_arr(1),
dtype=dtype)
self._gn3 = GroupNorm("gn3", out_filter, num_groups, dtype=dtype)
self._ds = stride[2] > 1 or in_filter != out_filter
if self._ds:
self._projconv = Conv2D("projconv",
1,
in_filter,
out_filter,
stride,
dtype=dtype)
self._projgn = GroupNorm("projgn",
out_filter,
num_groups,
dtype=dtype)
def __init__(self,
name,
layer_size,
add_bias=True,
bias_init=None,
act_func=None,
layernorm=False,
temp=None,
learn_temp=False,
dtype=tf.float32,
wdict=None):
super(CosineLastMLP, self).__init__(dtype=dtype)
self._layers = []
with variable_scope(name):
for i in range(len(layer_size) - 1):
if bias_init is not None and bias_init[i] is not None:
def bi():
return tf.zeros([layer_size[i + 1]],
dtype=dtype) + bias_init[i]
else:
bi = None
if i < len(layer_size) - 2:
layer = Linear("layer_{}".format(i),
layer_size[i],
layer_size[i + 1],
b_init=bi,
add_bias=add_bias,
dtype=dtype,
wdict=wdict)
else:
layer = CosineLinear("layer_{}".format(i),
layer_size[i],
layer_size[i + 1],
temp=temp,
learn_temp=learn_temp,
dtype=tf.float32,
wdict=wdict)
self._layers.append(layer)
if layernorm:
self._layers.append(
LayerNorm("layernorm_{}".format(i),
layer_size[i + 1],
dtype=dtype,
wdict=wdict))
if i < len(layer_size) - 2:
if act_func is None:
self._layers.append(tf.nn.relu)
else:
self._layers.append(act_func[i])
def __init__(self,
name,
in_filter,
out_filter,
stride,
num_groups=32,
data_format="NCHW",
dtype=tf.float32,
add_relu=True):
super(ResidualModule, self).__init__()
self._data_format = data_format
self._stride = stride
with variable_scope(name):
self._conv1 = Conv2D("conv1",
3,
in_filter,
out_filter,
stride,
data_format=data_format,
dtype=dtype)
self._gn1 = GroupNorm("gn1",
out_filter,
num_groups,
data_format=data_format,
dtype=dtype)
self._conv2 = Conv2D("conv2",
3,
out_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._gn2 = GroupNorm("gn2",
out_filter,
num_groups,
data_format=data_format,
dtype=dtype)
self._data_format = data_format
self._add_relu = True
self._ds = stride[2] > 1
self._pad = in_filter < out_filter
self._pad_size = [(out_filter - in_filter) // 2,
(out_filter - in_filter) // 2]
if data_format == "NCHW":
self._pad_arr = [[0, 0], self._pad_size, [0, 0], [0, 0]]
else:
self._pad_arr = [[0, 0], [0, 0], [0, 0], self._pad_size]
def __init__(self,
name,
in_filters,
num_hidden=384,
out_filters=512,
data_format="NCHW",
dtype=tf.float32):
super(FinalConvModule, self).__init__()
self._data_format = data_format
with variable_scope(name):
self._bn1 = BatchNorm("bn1", in_filters, data_format=data_format)
self._conv = Conv2D("conv", 1, in_filters, num_hidden,
self._stride_arr(1))
self._bn2 = BatchNorm("bn2", num_hidden, data_format=data_format)
self._fc = Linear("fc", num_hidden, out_filters)
self._out_filters = out_filters
def __init__(self, name, layer_size, dtype=tf.float32):
super(ResMLP, self).__init__(dtype=dtype)
self._layer_size = layer_size
self._layers = []
with variable_scope(name):
for i in range(len(layer_size) - 1):
def bi():
return tf.zeros([layer_size[i + 1]], dtype=dtype)
self._layers.append(
Linear("layer_{}".format(i),
layer_size[i],
layer_size[i + 1],
b_init=bi,
add_bias=True,
dtype=dtype))
def __init__(self, name, nin, nout, layernorm=False, dtype=tf.float32):
super(LSTM, self).__init__(dtype=dtype)
self._nin = nin
self._nout = nout
self._layernorm = layernorm
def _b_init():
return tf.concat(
[tf.ones([nout], dtype=dtype),
tf.zeros([3 * nout], dtype=dtype)],
axis=0)
with variable_scope(name):
self._gates = Linear(
"gates_linear", nin + nout, 4 * nout, b_init=_b_init)
if layernorm:
self._ln = LayerNorm("layernorm", 4 * nout, dtype=dtype)
def build_weight_init(self):
"""Initial state of the fast weights, with a mini batch.
Args:
bsize: Int. Batch size.
Returns:
fast_weights: List of fast weights, the first dimension is the batch
size.
"""
state_list = []
varlist = self.fast_weights
log.info('varlist {}'.format(varlist))
L = len(varlist)
shapelist = [[int(s) for s in v.shape] for v in varlist]
fast_weights = [None] * L
weight_init_list = []
with variable_scope(self._name):
for i in range(L):
key = self.fast_weights_keys[i]
if key.endswith('b'):
def binit():
return tf.zeros(shapelist[i])
b = self._get_variable('layer_{}/b_init'.format(i // 2),
binit,
dtype=self._dtype)
weight_init_list.append(b)
else:
def winit():
return tf.random.truncated_normal(shapelist[i],
mean=0.0,
stddev=0.01)
idx = i // 2 if self.config.classifier_bias else i
w = self._get_variable('layer_{}/w_init'.format(idx),
winit,
dtype=self._dtype)
weight_init_list.append(w)
return weight_init_list
def __init__(self,
name,
filter_size,
in_filters,
out_filters,
strides,
data_format='NCHW',
dtype=tf.float32,
wdict=None):
super(Conv2D, self).__init__(dtype=dtype)
with variable_scope(name):
kinit = self._get_normal_init(
[filter_size, filter_size, in_filters, out_filters])
self._kernel = self._get_variable("w",
kinit,
dtype=dtype,
wdict=wdict)
self._strides = strides
self._data_format = data_format
def __init__(self,
name,
num_channels,
data_format="NCHW",
eps=1e-3,
decay=0.999,
dtype=tf.float32,
wdict=None):
super(BatchNorm, self).__init__(dtype=dtype)
assert data_format in ["NCHW", "NHWC", "NHC", "NC"]
if data_format == "NCHW":
self._axis = 1
self._axes = [0, 2, 3]
elif data_format == "NHWC":
self._axis = -1
self._axes = [0, 1, 2]
elif data_format == "NC":
self._axis = -1
self._axes = [0]
elif data_format == "NHC":
self._axis = -1
self._axes = [0, 1]
self._eps = eps
self._decay = decay
self._data_format = data_format
C = num_channels
with variable_scope(name):
self._beta = self._get_variable("beta",
self._get_constant_init([C], 0.0),
wdict=wdict)
self._gamma = self._get_variable("gamma",
self._get_constant_init([C], 1.0),
wdict=wdict)
self._emean = self._get_variable("moving_mean",
self._get_constant_init([C], 0.0),
trainable=False,
wdict=wdict)
self._evar = self._get_variable("moving_variance",
self._get_constant_init([C], 0.0),
trainable=False,
wdict=wdict)
def __init__(self,
name,
in_dim,
out_dim,
w_init=None,
temp=None,
learn_temp=False,
dtype=tf.float32,
wdict=None):
super(CosineLinear, self).__init__(dtype=dtype)
if w_init is None:
w_init = self._get_uniform_init(in_dim, out_dim)
with variable_scope(name):
self._weight = self._get_variable("w", w_init, wdict=wdict)
self._in_dim = in_dim
self._out_dim = out_dim
self._name = name
if not learn_temp:
self._temp = temp
else:
self._temp = self._get_variable("temp", lambda: tf.zero([]) + temp)
def __init__(self,
name,
in_dim,
out_dim,
w_init=None,
b_init=None,
add_bias=True,
dtype=tf.float32,
wdict=None):
super(Linear, self).__init__(dtype=dtype)
if w_init is None:
w_init = self._get_uniform_init(in_dim, out_dim)
if b_init is None:
b_init = self._get_constant_init([out_dim], 0.0)
self._add_bias = add_bias
with variable_scope(name):
self._weight = self._get_variable("w", w_init, wdict=wdict)
if add_bias:
self._bias = self._get_variable("b", b_init, wdict=wdict)
self._in_dim = in_dim
self._out_dim = out_dim
self._name = name
def __init__(self,
name,
num_channels,
num_groups,
data_format="NCHW",
eps=1e-6,
dtype=tf.float32,
wdict=None):
super(GroupNorm, self).__init__(dtype=dtype)
if data_format == "NCHW":
channels_axis = 1
reduction_axes = (2, 3)
elif data_format == "NHWC":
channels_axis = 3
reduction_axes = (1, 2)
moment_axes = [channels_axis + 1]
for a in reduction_axes:
if a > channels_axis:
moment_axes.append(a + 1)
else:
moment_axes.append(a)
self._moment_axes = moment_axes
self._channels_axis = channels_axis
self._reduction_axes = reduction_axes
self._num_groups = num_groups
self._num_channels = num_channels
self._eps = eps
self._data_format = data_format
shape = [num_channels]
with variable_scope(name):
self._beta = self._get_variable("beta",
self._get_constant_init(
shape, 0.0),
wdict=wdict)
self._gamma = self._get_variable("gamma",
self._get_constant_init(
shape, 1.0),
wdict=wdict)
def __init__(self,
name,
nin,
nout,
nstack,
layernorm=False,
dtype=tf.float32):
super(StackLSTM, self).__init__(dtype=dtype)
self._nin = nin
self._nout = nout
self._lstm_list = []
self._nstack = nstack
assert nstack > 1, 'Number of LSTM > 1'
with variable_scope(name):
for n in range(nstack):
self._lstm_list.append(
LSTM(
"cell_{}".format(n),
nin,
nout,
layernorm=layernorm,
dtype=dtype))
nin = nout
def __init__(self,
name,
dim,
max_classes=20,
fix_unknown=False,
unknown_id=None,
similarity="euclidean",
temp_init=10.0,
dtype=tf.float32):
super(ProtoMemory, self).__init__(dtype=dtype)
self._max_classes = max_classes
self._fix_unknown = fix_unknown
self._unknown_id = unknown_id
self._similarity = similarity
self._dim = dim
if fix_unknown:
log.info('Fixing unknown id')
assert unknown_id is not None, 'Need to provide unknown ID'
if similarity in ["cosine", "poincare"]:
with variable_scope(name):
self._temp = self._get_variable("temp",
self._get_constant_init([], temp_init))
def __init__(self, name, layer_size, act_func, bias_init, dtype=tf.float32):
super(ResMLP, self).__init__(dtype=dtype)
self._layers = []
self._layer_size = layer_size
self._act_func = act_func
with variable_scope(name):
for i in range(len(layer_size) - 1):
if bias_init[i] is None:
bias_init_ = tf.zeros([layer_size[i + 1]], dtype=dtype)
else:
bias_init_ = bias_init[i]
def bi():
return bias_init_
self._layers.append(
Linear(
"layer_{}".format(i),
layer_size[i],
layer_size[i + 1],
b_init=bi,
add_bias=True,
dtype=dtype))
def __init__(self,
name,
filter_size,
in_filter,
out_filter,
leaky_relu=0.0,
max_pool=True,
data_format="NCHW",
dtype=tf.float32):
super(InitConvModule, self).__init__()
self._data_format = data_format
self._max_pool = max_pool
self._leaky_relu = leaky_relu
with variable_scope(name):
self._conv = Conv2D(
"conv",
filter_size,
in_filter,
out_filter,
self._stride_arr(1),
data_format=data_format,
dtype=dtype)
self._bn = BatchNorm(
"bn", out_filter, data_format=data_format, dtype=dtype)
def __init__(self,
name,
dim,
radius_init,
max_classes=20,
fix_unknown=False,
unknown_id=None,
similarity="euclidean",
static_beta_gamma=True,
radius_init_write=None,
use_ssl_beta_gamma_write=True,
dtype=tf.float32):
assert False, 'hey3'
super(SemiSupervisedMinDistProtoMemory,
self).__init__(name,
dim,
max_classes=max_classes,
fix_unknown=fix_unknown,
unknown_id=unknown_id,
similarity=similarity,
dtype=dtype)
self._controller_type = 'linear'
# self._controller_type = 'lstm'
self._radius_init = radius_init
log.info('Radius init {}'.format(radius_init))
if radius_init_write is not None:
self._radius_init_write = radius_init_write
log.info('Radius init write {}'.format(radius_init_write))
else:
self._radius_init_write = radius_init
self._use_ssl_beta_gamma_write = use_ssl_beta_gamma_write
if static_beta_gamma:
with variable_scope(name):
self._beta = self._get_variable(
"beta", self._get_constant_init([], radius_init))
self._gamma = self._get_variable(
"gamma", self._get_constant_init([], 1.0))
self._beta2 = self._get_variable(
"beta2",
self._get_constant_init([], self._radius_init_write))
self._gamma2 = self._get_variable(
"gamma2", self._get_constant_init([], 1.0))
with variable_scope(name):
if self._controller_type == 'lstm':
self._ctrl_lstm = LSTM("ctrl_lstm",
dim,
dim,
layernorm=False,
dtype=dtype)
self._ctrl_readout = Linear("ctrl_readout",
dim,
1,
w_init=lambda: tf.ones([dim, 1]),
b_init=lambda: tf.zeros([1]))
elif self._controller_type == 'linear':
self._ctrl_readout = Linear(
"ctrl_readout",
dim,
1,
# w_init=lambda: self._get_normal_init([dim, 1])() * 0.001,
w_init=lambda: tf.ones([dim, 1]) * 0.001,
b_init=lambda: tf.zeros([1]))
def __init__(self, config, wdict=None):
super(C4DoubleBackbone, self).__init__(config)
with variable_scope("b1"):
self.b1 = C4Backbone(config)
with variable_scope("b2"):
self.b2 = C4Backbone(config)
