def _test_flatten(test_case, device):
m = flow.nn.Flatten()
x = flow.Tensor(32, 2, 5, 5, device=flow.device(device))
flow.nn.init.uniform_(x)
y = m(x)
test_case.assertTrue(y.shape == flow.Size((32, 50)))
test_case.assertTrue(
np.array_equal(y.numpy().flatten(),
x.numpy().flatten()))
y2 = flow.flatten(x, start_dim=2)
test_case.assertTrue(y2.shape == flow.Size((32, 2, 25)))
test_case.assertTrue(
np.array_equal(y2.numpy().flatten(),
x.numpy().flatten()))
y3 = x.flatten(start_dim=1)
test_case.assertTrue(y3.shape == flow.Size((32, 50)))
test_case.assertTrue(
np.array_equal(y3.numpy().flatten(),
x.numpy().flatten()))
y4 = x.flatten(start_dim=1, end_dim=2)
test_case.assertTrue(y4.shape == flow.Size((32, 10, 5)))
test_case.assertTrue(
np.array_equal(y4.numpy().flatten(),
x.numpy().flatten()))
y5 = flow.flatten(x)
test_case.assertTrue(y5.shape == flow.Size((1600, )))
test_case.assertTrue(
np.array_equal(y5.numpy().flatten(),
x.numpy().flatten()))
def forward(self, x: flow.Tensor) -> flow.Tensor:
features = self.features(x)
out = F.relu(features, inplace=True)
out = F.adaptive_avg_pool2d(out, (1, 1))
out = flow.flatten(out, 1)
out = self.classifier(out)
return out
def _forward_impl(self, x: Tensor) -> Tensor:
if self.pad_input:
if self.channel_last:
# NHWC
paddings = (0, 1)
else:
# NCHW
paddings = (0, 0, 0, 0, 0, 1)
x = flow._C.pad(x, pad=paddings, mode="constant", value=0)
x = self.conv1(x)
if self.fuse_bn_relu:
x = self.bn1(x, None)
else:
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = flow.flatten(x, 1)
x = self.fc(x)
return x
def _forward_impl(self, x: Tensor) -> Tensor:
x = self.features(x)
# Cannot use "squeeze" as batch-size can be 1
x = self.adaptive_avg_pool2d(x)
x = flow.flatten(x, 1)
x = self.classifier(x)
return x
def alexnet(image, label, trainable=True):
conv1 = _conv2d_layer(
"conv1",
image,
filters=64,
kernel_size=11,
strides=4,
padding="VALID",
)
pool1 = flow.nn.avg_pool2d(conv1, 3, 2, "VALID", "NCHW", name="pool1")
conv2 = _conv2d_layer("conv2", pool1, filters=192, kernel_size=5)
pool2 = flow.nn.avg_pool2d(conv2, 3, 2, "VALID", "NCHW", name="pool2")
conv3 = _conv2d_layer("conv3", pool2, filters=384)
conv4 = _conv2d_layer("conv4", conv3, filters=384)
conv5 = _conv2d_layer("conv5", conv4, filters=256)
pool5 = flow.nn.avg_pool2d(conv5, 3, 2, "VALID", "NCHW", name="pool5")
if len(pool5.shape) > 2:
pool5 = flow.flatten(pool5, start_dim=1, end_dim=-1)
initializer = flow.truncated_normal_initializer(stddev=0.816496580927726)
fc1 = flow.layers.dense(
inputs=pool5,
units=4096,
activation=flow.math.relu,
use_bias=False,
kernel_initializer=initializer,
bias_initializer=False,
trainable=trainable,
name="fc1",
)
dropout1 = fc1
fc2 = flow.layers.dense(
inputs=dropout1,
units=4096,
activation=flow.math.relu,
use_bias=False,
kernel_initializer=initializer,
bias_initializer=False,
trainable=trainable,
name="fc2",
)
dropout2 = fc2
fc3 = flow.layers.dense(
inputs=dropout2,
units=1001,
activation=None,
use_bias=False,
kernel_initializer=initializer,
bias_initializer=False,
trainable=trainable,
name="fc3",
)
loss = flow.nn.sparse_softmax_cross_entropy_with_logits(
label, fc3, name="softmax_loss")
return loss
def SeModule(name, x, channel, reduction=4):
N, C, H, W = x.shape
y = flow.nn.avg_pool2d(x, ksize=[H, W], strides=None, padding="SAME")
y = flow.flatten(y, start_dim=1, end_dim=-1)
y = flow.layers.dense(
y,
units=channel // reduction,
use_bias=False,
kernel_initializer=_get_initializer("dense_weight"),
bias_initializer=_get_initializer("bias"),
kernel_regularizer=_get_regularizer("dense_weight"),
bias_regularizer=_get_regularizer("bias"),
name=name + "dense1a",
)
y = flow.math.relu(y)
y = flow.layers.dense(
y,
units=channel,
use_bias=False,
kernel_initializer=_get_initializer("dense_weight"),
bias_initializer=_get_initializer("bias"),
kernel_regularizer=_get_regularizer("dense_weight"),
bias_regularizer=_get_regularizer("bias"),
name=name + "dense2",
)
y = hsigmoid(y)
y = flow.expand_dims(input=y, axis=2)
y = flow.expand_dims(input=y, axis=3)
y_expand = flow.broadcast_like(y, x, broadcast_axes=(2, 3))
out = x * y_expand
return out
def forward(self, x):
x = self.pool(flow._C.relu(self.conv1(x)))
x = self.pool(flow._C.relu(self.conv2(x)))
x = flow.flatten(x, 1) # flatten all dimensions except batch
x = flow._C.relu(self.fc1(x))
x = flow._C.relu(self.fc2(x))
x = self.fc3(x)
return x
def forward(self):
train_record = self.train_record_reader()
label = self.record_label_decoder(train_record)
image_raw_buffer = self.record_image_decoder(train_record)
image = self.resize(image_raw_buffer)[0]
image = flow.flatten(image.to(flow.float32), start_dim=1)
return image, label
def forward(self, x):
# stage 0
x = flow.flatten(x, start_dim=1)
in0 = x.to_consistent(P0, B)
out0 = self.layer_0(in0)
# stage 1
in1 = out0.to_consistent(P1, B)
out1 = self.layer_1(in1)
return out1
def _forward_impl(self, x: Tensor) -> Tensor:
x = self.features(x)
x = self.avgpool(x)
x = flow.flatten(x, 1)
x = self.classifier(x)
return x
def forward(self, x):
out = self.stage0(x)
out = self.stage1(out)
out = self.stage2(out)
out = self.stage3(out)
out = self.stage4(out)
out = self.gap(out)
out = flow.flatten(out, 1)
out = self.linear(out)
return out
def forward(self, x: Tensor):
# N x 3 x 299 x 299
x = self.Conv2d_1a_3x3(x)
# N x 32 x 149 x 149
x = self.Conv2d_2a_3x3(x)
# N x 32 x 147 x 147
x = self.Conv2d_2b_3x3(x)
# N x 64 x 147 x 147
x = self.maxpool1(x)
# N x 64 x 73 x 73
x = self.Conv2d_3b_1x1(x)
# N x 80 x 73 x 73
x = self.Conv2d_4a_3x3(x)
# N x 192 x 71 x 71
x = self.maxpool2(x)
# N x 192 x 35 x 35
x = self.Mixed_5b(x)
# N x 256 x 35 x 35
x = self.Mixed_5c(x)
# N x 288 x 35 x 35
x = self.Mixed_5d(x)
# N x 288 x 35 x 35
x = self.Mixed_6a(x)
# N x 768 x 17 x 17
x = self.Mixed_6b(x)
# N x 768 x 17 x 17
x = self.Mixed_6c(x)
# N x 768 x 17 x 17
x = self.Mixed_6d(x)
# N x 768 x 17 x 17
x = self.Mixed_6e(x)
# N x 768 x 17 x 17
aux: Optional[Tensor] = None
if self.AuxLogits is not None:
if self.training:
aux = self.AuxLogits(x)
# N x 768 x 17 x 17
x = self.Mixed_7a(x)
# N x 1280 x 8 x 8
x = self.Mixed_7b(x)
# N x 2048 x 8 x 8
x = self.Mixed_7c(x)
# N x 2048 x 8 x 8
# Adaptive average pooling
x = self.avgpool(x)
# N x 2048 x 1 x 1
x = self.dropout(x)
# N x 2048 x 1 x 1
x = flow.flatten(x, 1)
# N x 2048
x = self.fc(x)
# N x 1000 (num_classes)
return x, aux
def _forward_impl(self, x: Tensor) -> Tensor:
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = flow.flatten(x, 1)
x = self.fc(x)
return x
def forward_features(self, x):
x = self.patch_embed(x)
if self.ape:
x = x + self.absolute_pos_embed
x = self.pos_drop(x)
for layer in self.layers:
x = layer(x)
x = self.norm(x) # B L C
x = self.avgpool(x.transpose(1, 2)) # B C 1
x = flow.flatten(x, 1)
return x
def forward(self, x: Tensor) -> Tensor:
# N x 768 x 17 x 17
x = self.avg_pool(x)
# N x 768 x 5 x 5
x = self.conv0(x)
# N x 128 x 5 x 5
x = self.conv1(x)
# N x 768 x 1 x 1
# Adaptive average pooling
x = self.adaptive_avp_pool(x)
# N x 768 x 1 x 1
x = flow.flatten(x, 1)
# N x 768
x = self.fc(x)
# N x 1000
return x
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.prelu(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.bn2(x)
x = flow.flatten(x, 1)
x = self.dropout(x)
x = self.fc(x)
x = self.features(x)
return x
def _forward_impl(self, x: Tensor) -> Tensor:
# See note [TorchScript super()]
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = flow.flatten(x, 1)
x = self.fc(x)
return x
def __init__(
self,
dim,
window_size,
num_heads,
qkv_bias=True,
qk_scale=None,
attn_drop=0.0,
proj_drop=0.0,
):
super().__init__()
self.dim = dim
self.window_size = window_size # Wh, Ww
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim**-0.5
# define a parameter table of relative position bias
# Author zzk: we add trunc normal here！
self.relative_position_bias_table = nn.Parameter(
flow.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1),
num_heads)) # 2*Wh-1 * 2*Ww-1, nH
self.relative_position_bias_table.trunc_normal_(std=0.02)
# get pair-wise relative position index for each token inside the window
coords_h = flow.arange(self.window_size[0])
coords_w = flow.arange(self.window_size[1])
coords = flow.stack(flow.meshgrid(*[coords_h, coords_w])) # 2, Wh, Ww
coords_flatten = flow.flatten(coords, 1) # 2, Wh*Ww
relative_coords = (coords_flatten[:, :, None] -
coords_flatten[:, None, :]) # 2, Wh*Ww, Wh*Ww
relative_coords = relative_coords.permute(1, 2, 0) # Wh*Ww, Wh*Ww, 2
relative_coords[:, :,
0] += self.window_size[0] - 1 # shift to start from 0
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
self.register_buffer("relative_position_index",
relative_position_index)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.softmax = nn.Softmax(dim=-1)
def _forward_impl(self, x: Tensor) -> Tuple[Any, Any]:
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
body = self.layer4(x)
x = self.avgpool(body)
x = flow.flatten(x, 1)
x = self.fc(x)
return x, body
def FlattenJob() -> flow.typing.Numpy:
with flow.scope.placement(device_type, "0:0"):
x = flow.get_variable(
"in",
shape=input_shape,
dtype=flow.float,
initializer=flow.random_uniform_initializer(minval=2, maxval=5),
trainable=True,
)
loss = flow.flatten(x, start_dim=start_dim, end_dim=end_dim)
flow.optimizer.SGD(
flow.optimizer.PiecewiseConstantScheduler([], [1e-4]), momentum=0
).minimize(loss)
flow.watch(x, test_global_storage.Setter("x"))
flow.watch_diff(x, test_global_storage.Setter("x_diff"))
return loss
def forward(self, x):
x = self.feature_extractor(x)
x = flow.flatten(x, 1)
logits = self.classifier(x)
probs = flow.softmax(logits, dim=1)
return logits, probs
def get_fc1(last_conv, num_classes, fc_type, input_channel=512):
body = last_conv
if fc_type == "Z":
body = _batch_norm(body,
epsilon=2e-5,
scale=False,
center=True,
is_training=True,
data_format="NCHW",
name="bn2")
body = _dropout(body, 0.4)
fc1 = body
elif fc_type == "E":
body = _batch_norm(body,
epsilon=2e-5,
is_training=True,
data_format="NCHW",
name="bn2")
body = _dropout(body, dropout_prob=0.4)
body = flow.flatten(body, 1)
fc1 = flow.layers.dense(
inputs=body,
units=num_classes,
activation=None,
use_bias=True,
kernel_initializer=_get_initializer(),
bias_initializer=flow.zeros_initializer(),
kernel_regularizer=_get_regularizer("weight"),
bias_regularizer=_get_regularizer("bias"),
trainable=True,
name="pre_fc1",
)
fc1 = _batch_norm(
fc1,
epsilon=2e-5,
scale=False,
center=True,
is_training=True,
data_format="NCHW",
name="fc1",
)
elif fc_type == "FC":
body = _batch_norm(body,
epsilon=2e-5,
is_training=True,
data_format="NCHW",
name="bn2")
body = flow.flatten(body, 1)
fc1 = flow.layers.dense(inputs=body,
units=num_classes,
activation=None,
use_bias=True,
kernel_initializer=_get_initializer(),
bias_initializer=flow.zeros_initializer(),
kernel_regularizer=_get_regularizer("weight"),
bias_regularizer=_get_regularizer("bias"),
trainable=True,
name="fc")
fc1 = _batch_norm(fc1,
epsilon=2e-5,
scale=False,
center=True,
is_training=True,
data_format="NCHW",
name="features")
elif fc_type == "GDC":
conv_6_dw = Linear(
last_conv,
num_filter=input_channel, # 512
num_group=input_channel, # 512
kernel=7,
pad="valid",
stride=[1, 1],
bn_is_training=True,
data_format="NCHW",
name="conv_6dw7_7",
)
conv_6_dw = flow.reshape(conv_6_dw, (body.shape[0], -1))
conv_6_f = flow.layers.dense(
inputs=conv_6_dw,
units=num_classes,
activation=None,
use_bias=True,
kernel_initializer=_get_initializer(),
bias_initializer=flow.zeros_initializer(),
kernel_regularizer=_get_regularizer("weight"),
bias_regularizer=_get_regularizer("bias"),
trainable=True,
name="pre_fc1",
)
fc1 = _batch_norm(
conv_6_f,
epsilon=2e-5,
scale=False,
center=True,
is_training=True,
data_format="NCHW",
name="fc1",
)
return fc1
def build_network(self,inputs):
b,c,t,h,w=inputs.shape
N=self.time_dim
templist=[]
for i in range(N):
tempname=datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S.%f')
if i!=N//2:
out = flow.range(t, dtype=flow.int64)
one = flow.constant_like(out, i, dtype= flow.int64)
out=flow.math.add(out, one)
out=flow.expand_dims(out,axis=0)
templist.append(out)
neighbor_time_index=flow.concat(templist,axis=0)
neighbor_time_index=flow.transpose(neighbor_time_index,[1,0])
neighbor_time_index=flow.flatten(neighbor_time_index, start_dim=0, end_dim=-1)
# feature map registration
tempname=datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S.%f')
init=flow.kaiming_initializer(shape=inputs.shape,mode="fan_out",nonlinearity="relu")
semantic=conv3d_layer("conv_semantic_"+tempname,inputs,self.out_channels,
kernel_size=1,use_bias=False,padding="VALID",trainable=self.trainable,
weight_initializer=init
)
inputs_norm=flow.math.l2_normalize(
semantic,axis=1
)
inputs_norm_padding=flow.pad(inputs_norm,paddings=[
(0,0),(0,0),((self.time_dim-1)//2,(self.time_dim-1)//2), (0,0),(0,0)]
)
inputs_norm_expand=flow.expand_dims(inputs_norm,axis=3)
temp_inputs_norm_expand=inputs_norm_expand
for i in range(N-2):
inputs_norm_expand=flow.concat(
inputs=[ inputs_norm_expand,temp_inputs_norm_expand],
axis=3
)
inputs_norm_expand=flow.transpose(inputs_norm_expand,perm=[0, 2, 3, 4, 5, 1])
inputs_norm_expand=flow.reshape(inputs_norm_expand,shape=[-1, h*w, c//16])
slice_list=[]
for index in neighbor_time_index:
temp=flow.slice(
inputs_norm_padding,
begin=[None,None,int(index),None,None],
size=[None,None,1,None,None]
)
slice_list.append(temp)
neighbor_norm=flow.concat(
slice_list,axis=2
)
neighbor_norm=flow.transpose(neighbor_norm,perm=[0, 2, 1, 3, 4])
neighbor_norm=flow.reshape(neighbor_norm,shape=[-1, c//16, h*w])
similarity=flow.matmul(inputs_norm_expand,neighbor_norm)*self.temperature
similarity=nn.softmax(similarity,axis=-1)
inputs_padding=flow.pad(inputs,
paddings=[
(0,0),(0,0),((self.time_dim-1)//2,(self.time_dim-1)//2), (0,0),(0,0)]
)
slice_list=[]
for index in neighbor_time_index:
temp=flow.slice(
inputs_padding,
begin=[None,None,int(index),None,None],
size=[None,None,1,None,None]
)
slice_list.append(temp)
neighbor=flow.concat(
slice_list,axis=2
)
neighbor=flow.transpose(neighbor,perm=[0,2,3,4,1])
neighbor=flow.reshape(neighbor,shape=[-1, h*w, c])
neighbor_new=flow.matmul(similarity,neighbor)
neighbor_new=flow.reshape(neighbor_new,shape=[b, t*(N-1), h, w, c])
neighbor_new=flow.transpose(neighbor_new,perm=[0, 4, 1, 2, 3])
# contrastive attention
if self.contrastive_att:
temp_input=flow.expand_dims(inputs,axis=3)
temp_temp_input=temp_input
for i in range(N-2):
temp_input=flow.concat(
inputs=[ temp_input,temp_temp_input],
axis=3
)
temp_input=flow.reshape(temp_input,shape=[b, c, (N-1)*t, h, w])
input_att=conv3d_layer(
"conv3d_inputmapping_"+tempname,temp_input,self.out_channels,
kernel_size=1, use_bias=False,trainable=False,weight_initializer=flow.kaiming_initializer(shape=temp_input.shape,mode="fan_out",nonlinearity="relu")
)
n_att=conv3d_layer(
"conv3d_nmapping_"+tempname,neighbor_new,self.out_channels,
kernel_size=1, use_bias=False,trainable=False,weight_initializer=flow.kaiming_initializer(shape=neighbor_new.shape,mode="fan_out",nonlinearity="relu")
)
temp_input=input_att*n_att
contrastive_att_net=conv3d_layer(
"conv3d_att_net_"+tempname,temp_input,1,
kernel_size=1, use_bias=False,trainable=self.trainable,weight_initializer=flow.kaiming_initializer(shape=temp_input.shape,mode="fan_out",nonlinearity="relu")
)
contrastive_att_net=flow.math.sigmoid(contrastive_att_net)
neighbor_new=flow.math.multiply(
neighbor_new,contrastive_att_net
)
# integrating feature maps
init = flow.zeros_initializer()
tempname=datetime.datetime.now().strftime('%Y-%m-%d-%H-%M-%S.%f')
input_offset = flow.get_variable(
"input_offset_"+tempname,
shape=(b, c, N*t, h, w),
initializer=init,
dtype=inputs.dtype,
trainable=self.trainable)
with flow.scope.placement("cpu", "0:0"):
input_index=np.array(
[i for i in range(t*N) if i%N==N//2]
)
neighbor_index=np.array(
[i for i in range(t*N) if i%N!=N//2])
input_offset_list=[]
inputs_list=[]
neighbor_new_list=[]
for index in range(input_offset.shape[2]):
temp=flow.slice(
input_offset,
begin=[None,None,int(index),None,None],
size=[None,None,1,None,None]
)
input_offset_list.append(temp)
for index in range(inputs.shape[2]):
temp=flow.slice(
inputs,
begin=[None,None,int(index),None,None],
size=[None,None,1,None,None]
)
inputs_list.append(temp)
for index in range(neighbor_new.shape[2]):
temp=flow.slice(
neighbor_new,
begin=[None,None,int(index),None,None],
size=[None,None,1,None,None]
)
neighbor_new_list.append(temp)
temp_index=0
for index in input_index:
input_offset_list[index]+=inputs_list[temp_index]
temp_index+=1
temp_index=0
for index in neighbor_index:
input_offset_list[index]+=neighbor_new_list[temp_index]
temp_index+=1
input_offset=flow.concat(
input_offset_list,axis=2
)
return input_offset
def forward(self, x: flow.Tensor) -> flow.Tensor:
x = self.features(x)
x = self.avgpool(x)
x = flow.flatten(x, 1)
x = self.classifier(x)
return x
def quantize_forward(self, x):
x = self.q_features(x)
x = self.q_avgpool(x)
x = flow.flatten(x, 1)
x = self.q_classifier(x)
return x
