def __init__(self, dim_in=4, dim_out=512, dim1=64, dim2=1024, im_size=64):
super(Encoder, self).__init__()
dim_hidden = [dim1, dim1 * 2, dim1 * 4, dim2, dim2]
self.conv1 = nn.Conv(dim_in,
dim_hidden[0],
kernel_size=5,
stride=2,
padding=2)
self.conv2 = nn.Conv(dim_hidden[0],
dim_hidden[1],
kernel_size=5,
stride=2,
padding=2)
self.conv3 = nn.Conv(dim_hidden[1],
dim_hidden[2],
kernel_size=5,
stride=2,
padding=2)
self.bn1 = nn.BatchNorm(dim_hidden[0])
self.bn2 = nn.BatchNorm(dim_hidden[1])
self.bn3 = nn.BatchNorm(dim_hidden[2])
self.fc1 = nn.Linear(dim_hidden[2] * math.ceil(im_size / 8)**2,
dim_hidden[3])
self.fc2 = nn.Linear(dim_hidden[3], dim_hidden[4])
self.fc3 = nn.Linear(dim_hidden[4], dim_out)
def __init__(self,
filename_obj,
dim_in=512,
centroid_scale=0.1,
bias_scale=1.0,
centroid_lr=0.1,
bias_lr=1.0):
super(Decoder, self).__init__()
# load .obj
self.template_mesh = jr.Mesh.from_obj(filename_obj)
# vertices_base, faces = jr.load_obj(filename_obj)
self.vertices_base = self.template_mesh.vertices[0].stop_grad(
) #vertices_base
self.faces = self.template_mesh.faces[0].stop_grad() #faces
self.nv = self.vertices_base.size(0)
self.nf = self.faces.size(0)
self.centroid_scale = centroid_scale
self.bias_scale = bias_scale
self.obj_scale = 0.5
dim = 1024
dim_hidden = [dim, dim * 2]
self.fc1 = nn.Linear(dim_in, dim_hidden[0])
self.fc2 = nn.Linear(dim_hidden[0], dim_hidden[1])
self.fc_centroid = nn.Linear(dim_hidden[1], 3)
self.fc_bias = nn.Linear(dim_hidden[1], self.nv * 3)
def __init__(self, n_classes=40):
super(DGCNN, self).__init__()
self.k = 20
self.knn = KNN(self.k)
self.bn1 = nn.BatchNorm(64)
self.bn2 = nn.BatchNorm(64)
self.bn3 = nn.BatchNorm(128)
self.bn4 = nn.BatchNorm(256)
self.bn5 = nn.BatchNorm1d(1024)
self.conv1 = nn.Sequential(nn.Conv(6, 64, kernel_size=1, bias=False),
self.bn1,
nn.LeakyReLU(scale=0.2))
self.conv2 = nn.Sequential(nn.Conv(64*2, 64, kernel_size=1, bias=False),
self.bn2,
nn.LeakyReLU(scale=0.2))
self.conv3 = nn.Sequential(nn.Conv(64*2, 128, kernel_size=1, bias=False),
self.bn3,
nn.LeakyReLU(scale=0.2))
self.conv4 = nn.Sequential(nn.Conv(128*2, 256, kernel_size=1, bias=False),
self.bn4,
nn.LeakyReLU(scale=0.2))
self.conv5 = nn.Sequential(nn.Conv1d(512, 1024, kernel_size=1, bias=False),
self.bn5,
nn.LeakyReLU(scale=0.2))
self.linear1 = nn.Linear(1024*2, 512, bias=False)
self.bn6 = nn.BatchNorm1d(512)
self.dp1 = nn.Dropout(p=0.5)
self.linear2 = nn.Linear(512, 256)
self.bn7 = nn.BatchNorm1d(256)
self.dp2 = nn.Dropout(p=0.5)
self.linear3 = nn.Linear(256, n_classes)
def __init__(self, dim_observation, dim_action):
super().__init__()
self.fc1 = nn.Linear(dim_observation + dim_action, 256)
# self.fc11 = nn.Linear(dim_observation, 256)
# self.fc12 = nn.Linear(dim_action, 256)
self.fc2 = nn.Linear(256, 256)
self.fc3 = nn.Linear(256, 1)
def __init__(self):
super(Discriminator, self).__init__()
def discriminator_block(in_filters, out_filters, bn=True):
'Returns layers of each discriminator block'
block = [
nn.Conv(in_filters, out_filters, 3, stride=2, padding=1),
nn.LeakyReLU(scale=0.2),
nn.Dropout(p=0.25)
]
if bn:
block.append(nn.BatchNorm(out_filters, eps=0.8))
return block
self.conv_blocks = nn.Sequential(
*discriminator_block(opt.channels, 16, bn=False),
*discriminator_block(16, 32), *discriminator_block(32, 64),
*discriminator_block(64, 128))
ds_size = (opt.img_size // (2**4))
self.adv_layer = nn.Sequential(nn.Linear((128 * (ds_size**2)), 1))
self.aux_layer = nn.Sequential(
nn.Linear((128 * (ds_size**2)), opt.n_classes), nn.Softmax())
self.latent_layer = nn.Sequential(
nn.Linear((128 * (ds_size**2)), opt.code_dim))
for m in self.modules():
weights_init_normal(m)
def __init__(self, in_channels, num_classes, conv_block=None):
super(InceptionAux, self).__init__()
if (conv_block is None):
conv_block = BasicConv2d
self.conv = conv_block(in_channels, 128, kernel_size=1)
self.fc1 = nn.Linear(2048, 1024)
self.fc2 = nn.Linear(1024, num_classes)
def __init__(self, n_classes=40):
super(PointConvDensityClsSsg, self).__init__()
self.sa1 = PointConvDensitySetAbstraction(npoint=512,
nsample=32,
in_channel=3,
mlp=[64, 64, 128],
bandwidth=0.1,
group_all=False)
self.sa2 = PointConvDensitySetAbstraction(npoint=128,
nsample=64,
in_channel=128 + 3,
mlp=[128, 128, 256],
bandwidth=0.2,
group_all=False)
self.sa3 = PointConvDensitySetAbstraction(npoint=1,
nsample=None,
in_channel=256 + 3,
mlp=[256, 512, 1024],
bandwidth=0.4,
group_all=True)
self.fc1 = nn.Linear(1024, 512)
self.bn1 = nn.BatchNorm1d(512)
self.drop1 = nn.Dropout(0.4)
self.fc2 = nn.Linear(512, 256)
self.bn2 = nn.BatchNorm1d(256)
self.drop2 = nn.Dropout(0.4)
self.fc3 = nn.Linear(256, n_classes)
self.relu = nn.ReLU()
def __init__(self, output_channels=40):
super(Point_Transformer, self).__init__()
self.conv1 = nn.Conv1d(3, 128, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(128, 128, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(128)
self.sa1 = SA_Layer(128)
self.sa2 = SA_Layer(128)
self.sa3 = SA_Layer(128)
self.sa4 = SA_Layer(128)
self.conv_fuse = nn.Sequential(
nn.Conv1d(512, 1024, kernel_size=1, bias=False),
nn.BatchNorm1d(1024), nn.LeakyReLU(scale=0.2))
self.linear1 = nn.Linear(1024, 512, bias=False)
self.bn6 = nn.BatchNorm1d(512)
self.dp1 = nn.Dropout(p=0.5)
self.linear2 = nn.Linear(512, 256)
self.bn7 = nn.BatchNorm1d(256)
self.dp2 = nn.Dropout(p=0.5)
self.linear3 = nn.Linear(256, output_channels)
self.relu = nn.ReLU()
def __init__(self):
super(SingleInputNet, self).__init__()
self.conv1 = nn.Conv(1, 10, 5)
self.conv2 = nn.Conv(10, 20, 5)
self.conv2_drop = nn.Dropout(p=0.3)
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def __init__(self, latent_dim, n_c, x_shape, verbose=False):
super(Generator_CNN, self).__init__()
self.name = 'generator'
self.latent_dim = latent_dim
self.n_c = n_c
self.x_shape = x_shape
self.ishape = (128, 7, 7)
self.iels = int(np.prod(self.ishape))
self.verbose = verbose
self.model0 = nn.Sequential(
nn.Linear((self.latent_dim + self.n_c), 1024))
self.model1 = nn.Sequential(BatchNorm1d(1024), nn.Leaky_relu(0.2))
self.model2 = nn.Sequential(nn.Linear(1024, self.iels),
BatchNorm1d(self.iels), nn.Leaky_relu(0.2))
self.model3 = nn.Sequential(
Reshape(self.ishape),
nn.ConvTranspose(128, 64, 4, stride=2, padding=1, bias=True),
nn.BatchNorm(64), nn.Leaky_relu(0.2))
self.model4 = nn.Sequential(
nn.ConvTranspose(64, 1, 4, stride=2, padding=1, bias=True))
self.sigmoid = nn.Sigmoid()
initialize_weights(self)
if self.verbose:
print('Setting up {}...\n'.format(self.name))
print(self.model)
def __init__(self):
super(Encoder, self).__init__()
self.model = nn.Sequential(nn.Linear(int(np.prod(img_shape)), 512),
nn.Leaky_relu(0.2), nn.Linear(512, 512),
nn.BatchNorm1d(512), nn.Leaky_relu(0.2))
self.mu = nn.Linear(512, opt.latent_dim)
self.logvar = nn.Linear(512, opt.latent_dim)
def __init__(self):
super(Decoder, self).__init__()
self.model = nn.Sequential(nn.Linear(opt.latent_dim, 512),
nn.Leaky_relu(0.2), nn.Linear(512, 512),
nn.BatchNorm1d(512), nn.Leaky_relu(0.2),
nn.Linear(512, int(np.prod(img_shape))),
nn.Tanh())
def __init__(self, dim, heads=8, dropout=0.):
super(Attention, self).__init__()
self.heads = heads
self.scale = dim**-0.5
self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
self.to_out = nn.Sequential(nn.Linear(dim, dim), nn.Dropout(dropout))
def __init__(self, cfg, in_channels):
super(MaskIoUFeatureExtractor, self).__init__()
layers = cfg.MODEL.ROI_MASKIOU_HEAD.CONV_LAYERS
resolution = cfg.MODEL.ROI_MASK_HEAD.POOLER_RESOLUTION // 2
input_features = in_channels + 1
fc_input_size = layers[0] * resolution * resolution
self.blocks = []
stride = 1
for layer_idx, layer_features in enumerate(layers, 1):
layer_name = "maskiou_fcn{}".format(layer_idx)
if layer_idx == len(layers):
stride = 2
module = make_conv3x3(input_features,
layer_features,
stride=stride)
setattr(self, layer_name, module)
input_features = layer_features
self.blocks.append(layer_name)
self.maskiou_fc1 = nn.Linear(fc_input_size, 1024)
self.maskiou_fc2 = nn.Linear(1024, 1024)
for l in [self.maskiou_fc1, self.maskiou_fc2]:
nn.init.kaiming_uniform_(l.weight, a=1)
nn.init.constant_(l.bias, 0)
def __init__(self):
super(Discriminator, self).__init__()
self.model = nn.Sequential(nn.Linear(int(np.prod(img_shape)),
512), nn.LeakyReLU(scale=0.2),
nn.Linear(512, 256),
nn.LeakyReLU(scale=0.2), nn.Linear(256, 1),
nn.Sigmoid())
def __init__(self):
super(Discriminator, self).__init__()
self.label_embedding = nn.Embedding(n_classes, n_classes)
self.model = nn.Sequential(
nn.Linear((n_classes + int(np.prod(img_shape))), 512),
nn.LeakyReLU(0.2), nn.Linear(512, 512), nn.Dropout(0.4),
nn.LeakyReLU(0.2), nn.Linear(512, 512), nn.Dropout(0.4),
nn.LeakyReLU(0.2), nn.Linear(512, 1))
def __init__(self, latent_dim, input_shape):
super(Encoder, self).__init__()
resnet18_model = resnet.Resnet18()
self.feature_extractor = nn.Sequential(*list(resnet18_model.children())[:(- 3)])
self.pooling = nn.Pool(kernel_size=8, stride=8, padding=0, op='mean')
self.fc_mu = nn.Linear(256, latent_dim)
self.fc_logvar = nn.Linear(256, latent_dim)
for m in self.modules():
weights_init_normal(m)
def __init__(self, channel, reduction=16):
super(SELayer, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(channel, channel // reduction, bias=False),
nn.Relu(),
nn.Linear(channel // reduction, channel, bias=False),
nn.Sigmoid()
)
def make_fc(dim_in, hidden_dim, use_gn=False):
if use_gn:
fc = nn.Linear(dim_in, hidden_dim, bias=False)
init.kaiming_uniform_(fc.weight, a=1)
return nn.Sequential(fc, group_norm(hidden_dim))
fc = nn.Linear(dim_in, hidden_dim)
init.kaiming_uniform_(fc.weight, a=1)
init.constant_(fc.bias, 0)
return fc
def __init__(self):
super(Model, self).__init__()
self.conv1 = nn.Conv(3, 32, 3, 1) # no padding
self.conv2 = nn.Conv(32, 64, 3, 1)
self.bn = nn.BatchNorm(64)
self.max_pool = nn.Pool(2, 2)
self.relu = nn.Relu()
self.fc1 = nn.Linear(64 * 12 * 12, 256)
self.fc2 = nn.Linear(256, 10)
def __init__(self, features, num_classes=1000, init_weights=True):
super(VGG, self).__init__()
self.features = features
self.classifier = nn.Sequential(
nn.Linear(512 * 7 * 7, 4096),
nn.ReLU(),
nn.Dropout(),
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Dropout(),
nn.Linear(4096, num_classes),
)
def __init__(
self,
embed_dim,
num_heads,
kdim=None,
vdim=None,
dropout=0.0,
bias=True,
add_bias_kv=False,
add_zero_attn=False,
self_attention=False,
encoder_decoder_attention=False,
q_noise=0.0,
qn_block_size=8,
):
super().__init__()
self.embed_dim = embed_dim
self.kdim = kdim if kdim is not None else embed_dim
self.vdim = vdim if vdim is not None else embed_dim
self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
self.num_heads = num_heads
assert dropout == 0, "TODO: dropout>0"
self.head_dim = embed_dim // num_heads
assert (self.head_dim * num_heads == self.embed_dim
), "embed_dim must be divisible by num_heads"
self.scaling = self.head_dim**-0.5
self.self_attention = self_attention
self.encoder_decoder_attention = encoder_decoder_attention
assert not self.self_attention or self.qkv_same_dim, (
"Self-attention requires query, key and "
"value to be of the same size")
#TODO: quant_noise
self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
assert not add_bias_kv, "TODO: add_bias_kv=True"
self.bias_k = self.bias_v = None
self.add_zero_attn = add_zero_attn
self.reset_parameters()
self.onnx_trace = False
self.tpu = False
def __init__(self,
in_features,
hidden_features=None,
out_features=None,
act_layer=nn.GELU,
drop=0.):
super(MLP, self).__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def __init__(self, cfg, in_channels):
super(FPNPredictor, self).__init__()
num_classes = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASSES
representation_size = in_channels
self.cls_score = nn.Linear(representation_size, num_classes)
num_bbox_reg_classes = 2 if cfg.MODEL.CLS_AGNOSTIC_BBOX_REG else num_classes
self.bbox_pred = nn.Linear(representation_size,
num_bbox_reg_classes * 4)
init.gauss_(self.cls_score.weight, std=0.01)
init.gauss_(self.bbox_pred.weight, std=0.001)
for l in [self.cls_score, self.bbox_pred]:
init.constant_(l.bias, 0)
def __init__(self):
super(Discriminator, self).__init__()
self.down = nn.Sequential(nn.Conv(opt.channels, 64, 3, stride=2, padding=1), nn.ReLU())
self.down_size = (opt.img_size // 2)
down_dim = (64 * ((opt.img_size // 2) ** 2))
self.embedding = nn.Linear(down_dim, 32)
self.fc = nn.Sequential(
nn.BatchNorm1d(32, 0.8),
nn.ReLU(),
nn.Linear(32, down_dim),
nn.BatchNorm1d(down_dim),
nn.ReLU()
)
self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv(64, opt.channels, 3, stride=1, padding=1))
def __init__(self, args, margs):
super(PointCloudTransformer, self).__init__(args, margs)
self.input_embeds = nn.Sequential(
Permute(0, 2, 1), nn.Conv1d(3, 64, kernel_size=1, bias=False),
nn.BatchNorm1d(64), nn.ReLU(),
nn.Conv1d(64, 64, kernel_size=1, bias=False), nn.BatchNorm1d(64),
nn.ReLU(), Permute(0, 2, 1))
self.knn_embeds = nn.Sequential(KNNEmbed(128, 128, 512, 32),
KNNEmbed(256, 256, 256, 32))
self.transformer = PointTransformer()
self.classifier = nn.Sequential(nn.Linear(1024, 512),
nn.BatchNorm1d(512), nn.ReLU(),
nn.Dropout(p=0.5), nn.Linear(512, 256),
nn.BatchNorm1d(256), nn.Dropout(p=0.5),
nn.Linear(256, 40))
def __init__(self):
super(STN3d, self).__init__()
self.conv1 = nn.Conv1d(3, 64, 1)
self.conv2 = nn.Conv1d(64, 128, 1)
self.conv3 = nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 9)
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
def __init__(self,
backbone,
img_size=224,
feature_size=None,
in_chans=3,
embed_dim=768):
super(HybridEmbed, self).__init__()
assert isinstance(backbone, nn.Module)
img_size = _pair(img_size)
self.img_size = img_size
self.backbone = backbone
if feature_size is None:
with jt.no_grad():
# FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
# map for all networks, the feature metadata has reliable channel and stride info, but using
# stride to calc feature dim requires info about padding of each stage that isn't captured.
training = backbone.is_training()
if training:
backbone.eval()
o = self.backbone(
jt.zeros((1, in_chans, img_size[0], img_size[1])))[-1]
feature_size = o.shape[-2:]
feature_dim = o.shape[1]
backbone.train()
else:
feature_size = _pair(feature_size)
feature_dim = self.backbone.feature_info.channels()[-1]
self.num_patches = feature_size[0] * feature_size[1]
self.proj = nn.Linear(feature_dim, embed_dim)
def __init__(self, alpha, num_classes=1000, dropout=0.2):
super(MNASNet, self).__init__()
assert (alpha > 0.0)
self.alpha = alpha
self.num_classes = num_classes
depths = _get_depths(alpha)
layers = [
nn.Conv(3, 32, 3, padding=1, stride=2, bias=False),
nn.BatchNorm(32, momentum=_BN_MOMENTUM),
nn.Relu(),
nn.Conv(32, 32, 3, padding=1, stride=1, groups=32, bias=False),
nn.BatchNorm(32, momentum=_BN_MOMENTUM),
nn.Relu(),
nn.Conv(32, 16, 1, padding=0, stride=1, bias=False),
nn.BatchNorm(16, momentum=_BN_MOMENTUM),
_stack(16, depths[0], 3, 2, 3, 3, _BN_MOMENTUM),
_stack(depths[0], depths[1], 5, 2, 3, 3, _BN_MOMENTUM),
_stack(depths[1], depths[2], 5, 2, 6, 3, _BN_MOMENTUM),
_stack(depths[2], depths[3], 3, 1, 6, 2, _BN_MOMENTUM),
_stack(depths[3], depths[4], 5, 2, 6, 4, _BN_MOMENTUM),
_stack(depths[4], depths[5], 3, 1, 6, 1, _BN_MOMENTUM),
nn.Conv(depths[5], 1280, 1, padding=0, stride=1, bias=False),
nn.BatchNorm(1280, momentum=_BN_MOMENTUM),
nn.Relu()
]
self.layers = nn.Sequential(*layers)
self.classifier = nn.Sequential(nn.Dropout(p=dropout),
nn.Linear(1280, num_classes))
def __init__(self, in_channel, style_dim):
super(AdaptiveInstanceNorm, self).__init__()
self.norm = nn.InstanceNorm2d(in_channel)
self.linear = nn.Linear(style_dim, in_channel * 2)
self.linear.bias.data[:in_channel] = 1
self.linear.bias.data[in_channel:] = 0
