def _forward(self, x):
x = self.Conv2d_1a_3x3(x)
x = self.Conv2d_2a_3x3(x)
x = self.Conv2d_2b_3x3(x)
x = nn.pool(x, 3, "maximum", stride=2)
x = self.Conv2d_3b_1x1(x)
x = self.Conv2d_4a_3x3(x)
x = nn.pool(x, 3, "maximum", stride=2)
x = self.Mixed_5b(x)
x = self.Mixed_5c(x)
x = self.Mixed_5d(x)
x = self.Mixed_6a(x)
x = self.Mixed_6b(x)
x = self.Mixed_6c(x)
x = self.Mixed_6d(x)
x = self.Mixed_6e(x)
aux_defined = self.aux_logits
if aux_defined:
aux = self.AuxLogits(x)
else:
aux = None
x = self.Mixed_7a(x)
x = self.Mixed_7b(x)
x = self.Mixed_7c(x)
x = nn.AdaptiveAvgPool2d(1)(x)
x = nn.Dropout()(x)
x = jt.reshape(x, (x.shape[0], (-1)))
x = self.fc(x)
return (x, aux)
def execute(self, inputs):
assert len(inputs) == len(self.in_channels)
outs = []
outs.append(inputs[0])
for i in range(1, len(inputs)):
outs.append(
nn.interpolate(inputs[i], scale_factor=2**i, mode='bilinear'))
out = jt.contrib.concat(outs, dim=1)
'''
if out.requires_grad and self.with_checkpoint:
out = checkpoint(self.reduction_conv, out)
else:
out = self.reduction_conv(out)
'''
out = self.reduction_conv(out)
outs = [out]
for i in range(1, self.num_level):
outs.append(
nn.pool(out, kernel_size=2**i, stride=2**i, op=self.pooling))
outputs = []
if self.share_conv:
for i in range(self.num_level):
outputs.append(self.fpn_conv(outs[i]))
else:
for i in range(self.num_level):
if not outs[i].is_stop_grad() and self.with_checkpoint:
tmp_out = checkpoint(self.fpn_conv[i], outs[i])
else:
tmp_out = self.fpn_conv[i](outs[i])
outputs.append(tmp_out)
return tuple(outputs)
def _forward(self, x):
branch3x3 = self.branch3x3(x)
branch3x3dbl = self.branch3x3dbl_1(x)
branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
branch_pool = nn.pool(x, 3, "maximum", stride=2)
outputs = [branch3x3, branch3x3dbl, branch_pool]
return outputs
def test_index_pool2(self):
pool = jt.nn.Pool(2, return_indices=True)
a = jt.array([1,0,0,1,
0,0,0,0,
0,0,0,0,
1,0,0,1]).reshape((1,1,4,4))
b, idx = pool(a)
assert (idx.data.reshape((4,)) == [0,1,2,3]).all()
def execute(self, x):
x = nn.pool(x, kernel_size=5, op="mean", stride=3)
x = self.conv0(x)
x = self.conv1(x)
x = nn.AdaptiveAvgPool2d(1)(x)
x = jt.reshape(x, (x.shape[0], (-1)))
x = self.fc(x)
return x
def execute(self, x, mask):
mask_pool = nn.pool(mask, kernel_size=2, stride=2, op='maximum')
x = jt.contrib.concat((x, mask_pool), 1)
for layer_name in self.blocks:
x = nn.relu(getattr(self, layer_name)(x))
x = x.reshape(x.shape[0], -1)
x = nn.relu(self.maskiou_fc1(x))
x = nn.relu(self.maskiou_fc2(x))
return x
def execute(self, convouts: List[jt.Var]):
"""
Args:
- convouts (list): A list of convouts for the corresponding layers in in_channels.
Returns:
- A list of FPN convouts in the same order as x with extra downsample layers if requested.
"""
out = []
x = jt.zeros((1, ))
for i in range(len(convouts)):
out.append(x)
# For backward compatability, the conv layers are stored in reverse but the input and output is
# given in the correct order. Thus, use j=-i-1 for the input and output and i for the conv layers.
j = len(convouts)
for lat_layer in self.lat_layers.layers.values():
j -= 1
if j < len(convouts) - 1:
_, _, h, w = convouts[j].shape
#print('hh',(h,w),x.shape[-2:])
x = nn.interpolate(x,
size=(h, w),
mode=self.interpolation_mode,
align_corners=False)
# x = interpolate(x, size=(h, w), mode=self.interpolation_mode, align_corners=False)
x = x + lat_layer(convouts[j])
out[j] = x
# This janky second loop is here because jtScript.
j = len(convouts)
for pred_layer in self.pred_layers.layers.values():
j -= 1
out[j] = pred_layer(out[j])
if self.relu_pred_layers:
out[j] = nn.relu(out[j])
cur_idx = len(out)
# In the original paper, this takes care of P6
if self.use_conv_downsample:
for downsample_layer in self.downsample_layers.layers.values():
out.append(downsample_layer(out[-1]))
else:
for idx in range(self.num_downsample):
# Note: this is an untested alternative to out.append(out[-1][:, :, ::2, ::2]). Thanks jtScript.
out.append(nn.pool(out[-1], 1, stride=2, op='maximum'))
if self.relu_downsample_layers:
for idx in range(len(out) - cur_idx):
out[idx] = nn.relu(out[idx + cur_idx])
return out
def _forward(self, x):
branch3x3 = self.branch3x3_1(x)
branch3x3 = self.branch3x3_2(branch3x3)
branch7x7x3 = self.branch7x7x3_1(x)
branch7x7x3 = self.branch7x7x3_2(branch7x7x3)
branch7x7x3 = self.branch7x7x3_3(branch7x7x3)
branch7x7x3 = self.branch7x7x3_4(branch7x7x3)
branch_pool = nn.pool(x, kernel_size=3, op="maximum", stride=2)
outputs = [branch3x3, branch7x7x3, branch_pool]
return outputs
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 = nn.pool(x, 3, "mean", stride=1, padding=1)
branch_pool = self.branch_pool(branch_pool)
outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
return outputs
def execute(self, x):
out = nn.relu(self.bn1(self.conv1(x)))
out = self.layer1(out)
out = self.layer2(out)
out = self.layer3(out)
out = self.layer4(out)
out = nn.pool(out, size=4, op="mean", padding=0)
out = jt.reshape(out, [out.shape[0], -1])
out = self.linear1(out)
out = self.linear2(out)
return out
def execute(self, mesh, mode=None):
image_size = self.image_size * (2 if self.anti_aliasing else 1)
images = srf.soft_rasterize(mesh.face_vertices, mesh.face_textures,
image_size, self.background_color,
self.near, self.far, self.fill_back,
self.eps, self.sigma_val, self.dist_func,
self.dist_eps, self.gamma_val,
self.aggr_func_rgb, self.aggr_func_alpha,
self.texture_type)
if self.anti_aliasing:
images = nn.pool(images, 2, "mean", stride=2)
return images
def _forward(self, x):
branch1x1 = self.branch1x1(x)
branch3x3 = self.branch3x3_1(x)
branch3x3 = [self.branch3x3_2a(branch3x3), self.branch3x3_2b(branch3x3)]
branch3x3 = jt.contrib.concat(branch3x3, dim=1)
branch3x3dbl = self.branch3x3dbl_1(x)
branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
branch3x3dbl = [self.branch3x3dbl_3a(branch3x3dbl), self.branch3x3dbl_3b(branch3x3dbl)]
branch3x3dbl = jt.contrib.concat(branch3x3dbl, dim=1)
branch_pool = nn.pool(x, kernel_size=3, op="mean", stride=1, padding=1)
branch_pool = self.branch_pool(branch_pool)
outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
return outputs
def test_unpool(self):
from jittor import nn
pool = nn.MaxPool2d(2, stride=2, return_indices=True)
unpool = nn.MaxUnpool2d(2, stride=2)
input = jt.array([[[[1., 2, 3, 4, 0], [5, 6, 7, 8, 0],
[9, 10, 11, 12, 0], [13, 14, 15, 16, 0],
[0, 0, 0, 0, 0]]]])
output, indices = pool(input)
out = unpool(output, indices, output_size=input.shape)
assert (out == jt.array([[[[0., 0., 0., 0., 0.], [0., 6., 0., 8., 0.],
[0., 0., 0., 0.,
0.], [0., 14., 0., 16., 0.],
[0., 0., 0., 0., 0.]]]])).all()
def _forward(self, x):
branch1x1 = self.branch1x1(x)
branch7x7 = self.branch7x7_1(x)
branch7x7 = self.branch7x7_2(branch7x7)
branch7x7 = self.branch7x7_3(branch7x7)
branch7x7dbl = self.branch7x7dbl_1(x)
branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
branch_pool = nn.pool(x, kernel_size=3, op="mean", stride=1, padding=1)
branch_pool = self.branch_pool(branch_pool)
outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
return outputs
def execute(self, x):
out = nn.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
# Squeeze
w = nn.pool(out, out.shape[2], 'maximum', 0)
w = nn.relu(self.fc1(w))
w = nn.Sigmoid()(self.fc2(w))
# Excitation
out = out * w # New broadcasting feature from v0.2!
out += self.shortcut(x)
out = nn.relu(out)
return out
def execute(self, x):
out = nn.relu(self.bn1(x))
shortcut = self.shortcut(out) if hasattr(self, 'shortcut') else x
out = self.conv1(out)
out = self.conv2(nn.relu(self.bn2(out)))
# Squeeze
w = nn.pool(out, out.shape[2], 'maximum', 0)
w = nn.relu(self.fc1(w))
w = self.act(self.fc2(w))
# Excitation
out = out * w
out += shortcut
return out
def resnet(x, is_train, block, layers, num_classes=1000):
layer_in_planes = 64
x = nn.conv(x, 3, layer_in_planes, 7, 3, 2)
x = nn.batch_norm(x, is_train)
x = nn.relu(x)
x = nn.pool(x, 3, "maximum", 1, 2)
x, layer_in_planes = block(x, is_train, 64, layers[0], layer_in_planes)
x, layer_in_planes = block(x, is_train, 128, layers[1], layer_in_planes, 2)
x, layer_in_planes = block(x, is_train, 256, layers[2], layer_in_planes, 2)
x, layer_in_planes = block(x, is_train, 512, layers[3], layer_in_planes, 2)
x = x.reindex_reduce("add", [x.shape[0], x.shape[1]],
["i0", "i1"]) / x.shape[2] / x.shape[3]
x = nn.linear(x, num_classes)
return x
def execute(self, x):
assert x.ndim == 4
# 2x2, float32 => downscale using _blur2d().
if self.blur is not None and x.dtype == jt.float32:
return self.blur(x)
# Apply gain.
if self.gain != 1:
x = x * self.gain
# No-op => early exit.
if self.factor == 1:
return x
# Large factor => downscale using tf.nn.avg_pool().
# NOTE: Requires tf_config['graph_options.place_pruned_graph']=True to work.
# return F.avg_pool2d(x, self.factor)
return nn.pool(x,
kernel_size=self.factor,
op='mean',
stride=self.factor)
def execute(self, input, step=0, alpha=-1):
for i in range(step, -1, -1):
index = self.n_layer - i - 1
if i == step:
out = self.from_rgb[index](input)
if i == 0:
out_std = jt.sqrt(out.var + 1e-8)
mean_std = out_std.mean()
mean_std = mean_std.expand(out.size(0), 1, 4, 4)
out = jt.cat([out, mean_std], 1)
out = self.progression[index](out)
if i > 0:
if i == step and 0 <= alpha < 1:
skip_rgb = nn.pool(input, 2)
skip_rgb = self.from_rgb[index + 1](skip_rgb)
out = (1 - alpha) * skip_rgb + alpha * out
out = out.squeeze(2).squeeze(2)
out = self.linear(out)
return out
def test_index_pool(self):
pool = jt.nn.Pool(2, return_indices=True)
a = jt.randn([10, 3, 100, 100])
b, idx = pool(a)
idx.sync()
def rasterize_rgbad(
faces,
textures=None,
image_size=DEFAULT_IMAGE_SIZE,
anti_aliasing=DEFAULT_ANTI_ALIASING,
near=DEFAULT_NEAR,
far=DEFAULT_FAR,
eps=DEFAULT_EPS,
background_color=DEFAULT_BACKGROUND_COLOR,
return_rgb=True,
return_alpha=True,
return_depth=True,
):
"""
Generate RGB, alpha channel, and depth images from faces and textures (for RGB).
Args:
faces (jittor.Var): Faces. The shape is [batch size, number of faces, 3 (vertices), 3 (XYZ)].
textures (jittor.Var): Textures.
The shape is [batch size, number of faces, texture size, texture size, texture size, 3 (RGB)].
image_size (int): Width and height of rendered images.
anti_aliasing (bool): do anti-aliasing by super-sampling.
near (float): nearest z-coordinate to draw.
far (float): farthest z-coordinate to draw.
eps (float): small epsilon for approximated differentiation.
background_color (tuple): background color of RGB images.
return_rgb (bool): generate RGB images or not.
return_alpha (bool): generate alpha channels or not.
return_depth (bool): generate depth images or not.
Returns:
dict:
{
'rgb': RGB images. The shape is [batch size, 3, image_size, image_size].
'alpha': Alpha channels. The shape is [batch size, image_size, image_size].
'depth': Depth images. The shape is [batch size, image_size, image_size].
}
"""
if textures is None:
inputs = [faces, None]
else:
inputs = [faces, textures]
if anti_aliasing:
# 2x super-sampling
rgb, alpha, depth = Rasterize(image_size * 2, near, far, eps,
background_color, return_rgb,
return_alpha, return_depth)(*inputs)
else:
rgb, alpha, depth = Rasterize(image_size, near, far, eps,
background_color, return_rgb,
return_alpha, return_depth)(*inputs)
# transpose & vertical flip
if return_rgb:
rgb = rgb.permute((0, 3, 1, 2))
# may need to look at this again because it seems to be very slow
rgb = rgb[:, :, list(reversed(range(rgb.shape[2]))), :]
if return_alpha:
alpha = alpha[:, list(reversed(range(alpha.shape[1]))), :]
if return_depth:
depth = depth[:, list(reversed(range(depth.shape[1]))), :]
if anti_aliasing:
# 0.5x down-sampling
if return_rgb:
rgb = nn.pool(rgb, 2, "mean", stride=2)
if return_alpha:
alpha = nn.pool(alpha.unsqueeze(1), 2, "mean", stride=2)
if return_depth:
depth = nn.pool(depth.unsqueeze(1), 2, "mean", stride=2)
ret = {
'rgb': rgb if return_rgb else None,
'alpha': alpha if return_alpha else None,
'depth': depth if return_depth else None,
}
return ret
def execute(self, x):
x = self.maskiou_net(x)
maskiou_p = nn.pool(x, kernel_size=x.shape[2],
op='maximum').squeeze(-1).squeeze(-1)
return maskiou_p
def execute(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = nn.relu(x)
x = nn.pool(x, op='maximum', kernel_size=3, stride=2, padding=1)
return x
def execute(self, x):
return [nn.pool(x, kernel_size=1, op="maximum", stride=2, padding=0)]
