def _get_mat3x3(self, image):
"""get perspective matrix used in the transformation
note: there are only 8 degrees of freedom in a perspective matrix, while the output matrix has 9 variables.
Args:
image (Tensor): input images (shape: n * 3 * 112 * 112)
Returns:
mat3x3 (Tensor): perspective matrix (shape: n * 3 * 3)
"""
x = self.stem(image)
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.fc(x)
s = self.input_size
# 0.01 here is a magic number. it aims to maintain identity transform at early stage of training
residual = x.reshape(-1, 3, 3) * 0.01
base = mge.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]]).astype("float32")
base = F.broadcast_to(base, residual.shape)
left_scale = mge.tensor([[s, 0, 0], [0, s, 0], [0, 0,
1]]).astype("float32")
left_scale = F.broadcast_to(left_scale, residual.shape)
right_scale = mge.tensor([[1 / s, 0, 0], [0, 1 / s, 0],
[0, 0, 1]]).astype("float32")
right_scale = F.broadcast_to(right_scale, residual.shape)
mat3x3 = F.matmul(left_scale, F.matmul(base + residual, right_scale))
return mat3x3
def forward(self, x):
x = self.stem(x)
x = self.features(x)
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.classifier(x)
return x
def forward(self, x):
features = self.features(x)
out = F.relu(features)
out = F.avg_pool2d(out, kernel_size=7,
stride=1).reshape(features.shape[0], -1)
out = self.classifier(out)
return out
def forward(self, x):
x = self.bn1(x)
x = self.dropout(x)
x = F.avg_pool2d(x, self.size)
x = F.flatten(x, 1)
x = self.fc(x)
x = self.bn2(x)
return x
def forward(self, x):
x = self.extract_features(x)["res3"]
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.fc(x)
return x
def forward(self, in_tensor):
avg_pool = F.avg_pool2d(in_tensor,
(in_tensor.shape[2], in_tensor.shape[3]),
stride=(in_tensor.shape[2],
in_tensor.shape[3]))
x = self.gate_c(avg_pool)
x = F.expand_dims(x, axis=[2, 3]) # b,48,1
x = F.broadcast_to(x, in_tensor.shape) # b,48,h,w
return x
def forward(self, x):
x = F.reshape(x, (1, 3, 224, 224))
x = self.extract_features(x)["res5"]
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.fc(x)
return x
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = F.relu(x)
x = self.avgpool(x)
x = F.avg_pool2d(x, 22)
x = F.flatten(x, 1)
x = self.fc(x)
return x
def forward(self, x):
x = self.extract_features(x)["res5"]
x = F.avg_pool2d(x, 7)
print(x.shape)
x = F.flatten(x, 1)
x = self.fc(x)
return x
def _shortcut(self, x):
"""
Helper function for feedforwarding through shortcut layers.
"""
if self.learnable_sc:
x = self.c_sc(x)
return F.avg_pool2d(x, 2) if self.downsample else x
else:
return x
def _residual(self, x):
"""
Helper function for feedforwarding through main layers.
"""
h = x
h = self.c1(h)
h = self.activation(h)
h = self.c2(h)
h = F.avg_pool2d(h, 2)
return h
def forward(self, x):
x = self.quant(x)
x = self.first_conv(x)
x = self.maxpool(x)
x = self.features(x)
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.dequant(x)
x = self.classifier(x)
return x
def forward(self, x):
if dist.get_rank() > 0:
x = recv_fr_prev_gpu()
x = self.features(x)
if dist.get_rank() != 3:
_ = send_to_next_gpu(x)
else:
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.classifier(x)
return x
def forward(self, x):
# FIXME whenever finding elegant solution
x = self.quant(x)
x = self.extract_features(x)["res5"]
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.dequant(x)
x = self.fc(x)
return x
def _residual(self, x):
h = x
h = layernorm(h)
h = self.activation(h)
h = self.c1(h)
h = layernorm(h)
h = self.activation(h)
h = self.c2(h)
if self.downsample:
h = F.avg_pool2d(h, 2)
return h
def forward(self, tenFirst, tenSecond):
tenFirst = [self.preprocess(tenFirst)]
tenSecond = [self.preprocess(tenSecond)]
for intLevel in range(self.num_layers - 1):
if tenFirst[0].shape[2] >= self.threshold or tenFirst[0].shape[
3] >= self.threshold:
tenFirst.insert(
0, F.avg_pool2d(inp=tenFirst[0], kernel_size=2, stride=2))
tenSecond.insert(
0, F.avg_pool2d(inp=tenSecond[0], kernel_size=2, stride=2))
tenFlow = F.zeros([
tenFirst[0].shape[0], 2,
int(math.floor(tenFirst[0].shape[2] / 2.0)),
int(math.floor(tenFirst[0].shape[3] / 2.0))
])
# print(len(tenFirst))
for intLevel in range(len(tenFirst)):
# normal: 5 for training (4*4, 8*8, 16*16, 32*32, 64*64) 5 for test (11*20, 22*40, 45*80, 90*160, 180*320)
# small: 3 for training (16*16, 32*32, 64*64) 3 for test (45*80, 90*160, 180*320)
tenUpsampled = F.nn.interpolate(inp=tenFlow,
scale_factor=2,
mode='BILINEAR',
align_corners=True) * 2.0
if tenUpsampled.shape[2] != tenFirst[intLevel].shape[2]:
tenUpsampled = pad_H(tenUpsampled)
if tenUpsampled.shape[3] != tenFirst[intLevel].shape[3]:
tenUpsampled = pad_W(tenUpsampled)
tenFlow = self.netBasic[intLevel](F.concat([
tenFirst[intLevel],
backwarp(tenInput=tenSecond[intLevel],
tenFlow=tenUpsampled,
border_mode=self.border_mode), tenUpsampled
],
axis=1)) + tenUpsampled
return tenFlow
def forward(self, x):
branch1x1 = self.branch1x1(x)
branch5x5 = self.branch5x5_1(x)
branch5x5 = self.branch5x5_2(branch5x5)
branch3x3dbl = self.branch3x3dbl_1(x)
branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
branch_pool = self.branch_pool(branch_pool)
outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
return F.concat(outputs, 1)
def forward(self, inps):
if self._opr.name == "MaxPool2d":
return F.max_pool2d(
inps[0],
kernel_size=self.param["kernel_size"],
stride=self.param["stride"],
padding=self.param["padding"],
)
else:
return F.avg_pool2d(
inps[0],
kernel_size=self.param["kernel_size"],
stride=self.param["stride"],
padding=self.param["padding"],
mode=self.param["mode"],
)
def forward(self, inputs, _ref=None):
avgout = self.avg_pool(inputs)
eesp_out = self.eesp(inputs)
net = F.concat([eesp_out, avgout], axis=1)
if self.refin:
w1 = avgout.shape[2]
w2 = _ref.shape[2]
while w2 != w1:
_ref = F.avg_pool2d(_ref,
kernel_size=3,
stride=self.stride,
padding=1)
w2 = _ref.shape[2]
_ref = self.refin_conv(_ref)
net = net + _ref
net = self.activation(net)
return net
def forward(self, x):
self.num_chunks = 4
self.inp_chunks = []
self.oup_chunks = []
if dist.get_rank() == 0:
self.inp_chunks = F.split(x, 4)
for i in range(self.num_chunks):
if dist.get_rank() == 0:
x = self.inp_chunks[i]
else:
x = recv_fr_prev_gpu()
self.inp_chunks.append(x)
x = self.features(x)
if dist.get_rank() != 3:
_ = send_to_next_gpu(x)
else:
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
x = self.classifier(x)
self.oup_chunks.append(x)
return F.concat(self.oup_chunks)
def backward(self, label, gm):
label_chunks = F.split(label, 4)
losses = []
for i, x in enumerate(self.inp_chunks):
with gm: #ad.GradManager().attach(self.parameters()) as gm:
gm.attach(x) # query gradient of the input
y = self.features(x)
if dist.get_rank() == 3:
y = F.avg_pool2d(y, 7)
y = F.flatten(y, 1)
y = self.classifier(y)
loss = F.nn.cross_entropy(y, label_chunks[i])
losses.append(loss)
gm.backward(loss)
else:
grad = grad_fr_next_gpu()
gm.backward(y, dy=grad)
if dist.get_rank() != 0:
_ = grad_to_prev_gpu(x.grad)
return sum(losses) / self.num_chunks if losses else None
def forward(self, x):
h = x.shape[-2]
w = x.shape[-1]
x = F.avg_pool2d(x, [h, w])
return x
def get_stats(self, x):
flops, activations = 0, 0
in_x = deepcopy(x)
x = self.conv1(x)
tmp_flops, tmp_acts = cal_conv_stats(self.conv1, in_x, x)
activations += tmp_acts
flops += tmp_flops
in_x = deepcopy(x)
x = self.bn1(x)
tmp_flops, tmp_acts = cal_norm_stats(self.bn1, in_x, x)
activations += tmp_acts
flops += tmp_flops
x = F.relu(x)
in_x = deepcopy(x)
x = self.maxpool(x)
tmp_flops, tmp_acts = cal_pool_stats(self.maxpool, in_x, x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer1_0.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer1_1.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer2_0.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer2_1.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer3_0.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer3_1.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer4_0.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x, tmp_flops, tmp_acts = self.layer4_1.get_stats(x)
activations += tmp_acts
flops += tmp_flops
x = F.avg_pool2d(x, 7)
x = F.flatten(x, 1)
in_x = deepcopy(x)
x = self.fc(x)
tmp_flops, tmp_acts = cal_linear_stats(self.fc, in_x, x)
activations += tmp_acts
flops += tmp_flops
return flops, activations
def fwd(data):
out = F.max_pool2d(data, 2, 2)
out = F.avg_pool2d(out, 2, 2)
return out
def _avg_pool3x3(x):
xx = F.avg_pool2d(x, kernel_size=3, stride=1)
return xx
),
(
"adaptive_max_pool2d",
lambda x: MF.adaptive_max_pool2d(x, (7, 7)),
lambda x: TF.adaptive_max_pool2d(x, (7, 7)),
[(2, 32, 16, 16)],
[(64, 512, 16, 16)],
True,
1000,
),
("argsort", MF.argsort, torch.argsort, [(1000, )], [
(1000, 1000),
], True, 1000),
(
"avg_pool2d",
lambda x: MF.avg_pool2d(x, 2),
lambda x: TF.avg_pool2d(x, 2),
[(2, 32, 16, 16)],
[(64, 512, 16, 16)],
True,
1000,
),
(
"broadcast",
lambda x: MF.broadcast_to(x, (5, ) + x.shape),
lambda x: torch.broadcast_to(x, (5, ) + x.shape),
[(100, 100)],
[(64, 512, 16, 16)],
True,
1000,
),
def _shortcut(self, x):
"""
Helper function for feedforwarding through shortcut layers.
"""
return self.c_sc(F.avg_pool2d(x, 2))
