def __init__(self, mode='combined', activate=None): super(SegNet, self).__init__() _, self.data_type = get_device() self.mode = mode self.activate = activate self.dropout = nn.Dropout2d(0.3) self.maxpool2d = nn.MaxPool2d(2) self.upsample = nn.UpsamplingBilinear2d(scale_factor=2) self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=True) self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=True) self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1, bias=True) self.conv4 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1, bias=True) self.deconv1 = nn.ConvTranspose2d(512, 256, kernel_size=3, stride=1, padding=1, bias=True) self.deconv2 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=1, padding=1, bias=True) self.deconv3 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=1, padding=1, bias=True) self.deconv4 = nn.ConvTranspose2d(64, 3, kernel_size=3, stride=1, padding=1, bias=True) self.batch_norm = nn.BatchNorm2d(3) self.threshold1 = nn.Threshold(0.25, 0) self.threshold2 = nn.Threshold(0.5, 0) self.threshold3 = nn.Threshold(0.75, 0) self.maxpool1d = nn.MaxPool1d(3, stride=2, return_indices=True)
def __init__(self, config, embedding): """ :param hidden: BERT model output size """ super().__init__() self.problem = config.problem self.embedding_size = config.embedding_size self.batch_size = config.batch_size self.bert = BERT(config, embedding) self.model_type = config.model_type self.tau = config.tau self.norm = LayerNorm(self.embedding_size) self.linear_final = None if self.model_type in ['bert_max_kernel', 'bert_avg_kernel']: self.linear_final = nn.Linear(config.embedding_size * 5, config.n_classes) elif self.model_type == 'bert_max_min': self.linear_final = nn.Linear(config.embedding_size * 2, config.n_classes) elif self.model_type.find('bert_rnn') != -1: self.rnn = RNNEncoder(config) self.linear_final = nn.Linear(config.hidden_size, config.n_classes) elif self.model_type == 'bert_cnn': self.cnn = CNNEncoder(config) self.linear_final = nn.Linear( len(config.kernel_heights) * config.out_channels, config.n_classes) elif self.model_type == 'bert_sum': # self.norm = LayerNorm(config.embedding_size) # self.linear_dense = nn.Linear(config.embedding_size, config.embedding_size // 2) # self.linear_final = nn.Linear(config.embedding_size // 2, config.n_classes) self.linear_final = nn.Linear(config.embedding_size, config.n_classes) elif self.model_type == 'bert_weight': # Transformer_Adaptive self.weight = nn.Parameter( torch.zeros(config.batch_size, config.max_len, 1, device=config.device)) self.weight.data.fill_(0) # self.weight.data.uniform_(-np.sqrt(3.0 / config.embedding_size), np.sqrt(3.0 / config.embedding_size)) elif self.model_type == 'bert_weight_embedding': self.weight = nn.Parameter( torch.zeros(config.batch_size, config.embedding_size, 1, device=config.device)) # self.weight.data.fill_(0) self.weight.data.uniform_(-np.sqrt(3.0 / config.embedding_size), np.sqrt(3.0 / config.embedding_size)) elif self.model_type == 'bert_sample_manual': self.sample_weight = torch.ones(config.max_len) / config.max_len self.sample_num = 6 self.sample_weight[:20] = 0.1 self.sample_weight[:10] = 0.2 self.sample_weight[:5] = 0.3 self.linear_final = nn.Linear( config.embedding_size * self.sample_num, config.n_classes) elif self.model_type == 'bert_sample_exp': self.linear_final = nn.Linear(config.embedding_size, config.n_classes) elif self.model_type.startswith('bert_gumbel'): self.linear_final = nn.Linear(config.embedding_size, config.n_classes) elif self.model_type == 'bert_conv1d': self.conv1d1 = nn.Conv1d(config.embedding_size, config.embedding_size, 6) # 45 self.max_pool1d1 = nn.MaxPool1d(3) # 15 self.conv1d2 = nn.Conv1d(config.embedding_size, config.embedding_size, 4) # 12 self.max_pool1d2 = nn.MaxPool1d(2) # 6 self.conv1d3 = nn.Conv1d(config.embedding_size, config.embedding_size, 4) # 3 self.max_pool1d3 = nn.MaxPool1d(3) # 1 # self.linear_dense = nn.Linear(config.embedding_size, config.embedding_size // 2) # self.dropout_dense = nn.Dropout(0.5) # self.linear_final = nn.Linear(config.embedding_size // 2, config.n_classes) # self.norm = LayerNorm(config.embedding_size) self.linear_final = nn.Linear(config.embedding_size, config.n_classes) if self.linear_final is None: if self.model_type in [ 'bert_max_embedding', 'bert_avg_embedding', 'bert_weight_embedding' ]: self.linear_final = nn.Linear(config.max_len, config.n_classes) else: self.linear_final = nn.Linear(config.embedding_size, config.n_classes) if self.model_type.startswith( 'bert_weight') or self.model_type.startswith('bert_gumbel'): self.position_weights = None
maxpooling_sigmoid = nn.Sequential( nn.MaxPool2d(kernel_size=4, stride=2, padding=(1, 2), dilation=1), nn.Sigmoid() ) input = Variable(torch.randn(2, 3, 12, 18)) save_data_and_model("maxpooling_sigmoid_dynamic_axes", input, maxpooling_sigmoid) postprocess_model("models/maxpooling_sigmoid_dynamic_axes.onnx", [[2, 3, 'height', 'width']]) ave_pool = nn.AvgPool2d(kernel_size=3, stride=2, padding=1) input = Variable(torch.randn(1, 3, 7, 5)) save_data_and_model("average_pooling_dynamic_axes", input, ave_pool) postprocess_model("models/average_pooling_dynamic_axes.onnx", [[1, 3, 'height', 'width']]) x = Variable(torch.randn(1, 3, 10)) max_pool = nn.MaxPool1d(kernel_size=(5), stride=1, padding=2, dilation=1) save_data_and_model("maxpooling_1d", x, max_pool) x = Variable(torch.randn(2, 3, 12)) maxpooling_sigmoid = nn.Sequential( nn.MaxPool1d(kernel_size=4, stride=2, padding=(2), dilation=1), nn.Sigmoid() ) save_data_and_model("maxpooling_sigmoid_1d", x, maxpooling_sigmoid) x = Variable(torch.randn(2, 3, 12)) maxpool2 = nn.Sequential( nn.MaxPool1d(kernel_size=5, stride=1, padding=0, dilation=1), nn.MaxPool1d(kernel_size=3, stride=1, padding=0, dilation=1) ) save_data_and_model("two_maxpooling_1d", x, maxpool2)
def __init__(self, in_channels, inter_channels, scale, dimension=3, sub_sample=True, bn_layer=True, in_layer=False): super(Adaptive_NL2_BlockND, self).__init__() assert dimension in [1, 2, 3] self.dimension = dimension self.sub_sample = sub_sample self.in_channels = in_channels self.inter_channels = inter_channels self.scale = scale # multi-scale,可选scale if dimension==3: conv_nd = nn.Conv3d max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2)) self.adaptive_pool_layer = nn.AdaptiveAvgPool3d((1, scale, scale)) # 用于第二个分支的Adaptive Pooling bn = nn.BatchNorm3d instanceNorm = nn.InstanceNorm3d # 实例归一化 elif dimension == 2: conv_nd = nn.Conv2d max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2)) self.adaptive_pool_layer = nn.AdaptiveAvgPool2d((scale, scale)) # 用于第二个分支的Adaptive Pooling bn = nn.BatchNorm2d instanceNorm = nn.InstanceNorm2d # 实例归一化 else: conv_nd = nn.Conv1d max_pool_layer = nn.MaxPool1d(kernel_size=2) self.adaptive_pool_layer = nn.AdaptiveAvgPool1d(scale) # 用于第二个分支的Adaptive Pooling bn = nn.BatchNorm1d instanceNorm = nn.InstanceNorm1d # 实例归一化 if bn_layer: self.W = nn.Sequential( conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0), bn(self.in_channels) ) nn.init.constant_(self.W[1].weight, 0) # 只对bn参数初始化,因为nn.conv2d会自动初始化参数 nn.init.constant_(self.W[1].bias, 0) elif in_layer: self.W = nn.Sequential( conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, # channel # 减半,减少计算量: 1024 ->512 kernel_size=1, stride=1, padding=0), instanceNorm(self.in_channels) ) nn.init.constant_(self.W[1].weight, 0) # 只对bn参数初始化,因为nn.conv2d会自动初始化参数 nn.init.constant_(self.W[1].bias, 0) else: self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0) nn.init.constant_(self.W.weight, 0) nn.init.constant_(self.W.bias, 0) self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0) self.phi = conv_nd(in_channels=self.in_channels, out_channels=1, kernel_size=1, stride=1, padding=0) self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0) self.alpha = conv_nd(in_channels=self.inter_channels, out_channels=self.scale*self.scale, kernel_size=1, stride=1, padding=0) self.beta = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0) if sub_sample: self.g = nn.Sequential(self.g, max_pool_layer) self.phi = nn.Sequential(self.phi, max_pool_layer)
def test_maxpool_dilations(self): x = torch.randn(20, 16, 50) self.assertONNX(nn.MaxPool1d(2, stride=1, dilation=2), x, opset_version=10)
def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True): super(_NonLocalBlockND, self).__init__() assert dimension in [1, 2, 3] self.dimension = dimension self.sub_sample = sub_sample self.in_channels = in_channels self.inter_channels = inter_channels if self.inter_channels is None: self.inter_channels = in_channels // 2 if self.inter_channels == 0: self.inter_channels = 1 if dimension == 3: conv_nd = nn.Conv3d max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2)) bn = nn.BatchNorm3d elif dimension == 2: conv_nd = nn.Conv2d max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2)) bn = nn.BatchNorm2d else: conv_nd = nn.Conv1d max_pool_layer = nn.MaxPool1d(kernel_size=(2)) bn = nn.BatchNorm1d self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels, kernel_size=1, stride=1, padding=0) if bn_layer: self.W = nn.Sequential( conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0), bn(self.in_channels)) nn.init.constant(self.W[1].weight, 0) nn.init.constant(self.W[1].bias, 0) else: self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels, kernel_size=1, stride=1, padding=0) nn.init.constant(self.W.weight, 0) nn.init.constant(self.W.bias, 0) if sub_sample: self.g = nn.Sequential(self.g, max_pool_layer) self.phi = max_pool_layer
def __init__(self, input_channel, layers=[1, 1, 1, 1], num_classes=10): self.inplanes3_1 = 64 self.inplanes3_2 = 64 self.inplanes3_3 = 64 super(MSResNet, self).__init__() self.conv1 = nn.Conv1d(input_channel, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm1d(64) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1) self.layer3x3_11 = self._make_layer3_1(BasicBlock3x3_1, 64, layers[0], stride=2) self.layer3x3_12 = self._make_layer3_1(BasicBlock3x3_1, 128, layers[1], stride=2) self.layer3x3_13 = self._make_layer3_1(BasicBlock3x3_1, 256, layers[2], stride=2) # self.layer3x3_4 = self._make_layer3(BasicBlock3x3, 512, layers[3], stride=2) # maxplooing kernel size: 16, 11, 6 self.maxpool3_1 = nn.AvgPool1d(kernel_size=16, stride=1, padding=0) self.layer3x3_21 = self._make_layer3_2(BasicBlock3x3_2, 64, layers[0], stride=2) self.layer3x3_22 = self._make_layer3_2(BasicBlock3x3_2, 128, layers[1], stride=2) self.layer3x3_23 = self._make_layer3_2(BasicBlock3x3_2, 256, layers[2], stride=2) # self.layer3x3_4 = self._make_layer3(BasicBlock3x3, 512, layers[3], stride=2) # maxplooing kernel size: 16, 11, 6 self.maxpool3_2 = nn.AvgPool1d(kernel_size=16, stride=1, padding=0) self.layer3x3_31 = self._make_layer3_3(BasicBlock3x3_3, 64, layers[0], stride=2) self.layer3x3_32 = self._make_layer3_3(BasicBlock3x3_3, 128, layers[1], stride=2) self.layer3x3_33 = self._make_layer3_3(BasicBlock3x3_3, 256, layers[2], stride=2) # self.layer3x3_4 = self._make_layer3(BasicBlock3x3, 512, layers[3], stride=2) # maxplooing kernel size: 16, 11, 6 self.maxpool3_3 = nn.AvgPool1d(kernel_size=16, stride=1, padding=0) # self.drop = nn.Dropout(p=0.2) self.fc = nn.Linear(256 * 3, num_classes)
def __init__(self, in_channels, seq_len, out_channels, cnnfilter_size, cnnfilter_stride, cnn_pad, use_bias, use_maxpool, pool_size, pool_stride, use_batchnorm, eps=1e-5, momentum=0.9, affine=False, dropout=0): """ For Building Repeatable Blocks. Just CNN_1d, ReLU, Opt. Batchnorm, Opt. Dropout, Opt. MaxPool in_channels = number of input features (like color channels in RBG images for example) seq_len = length of our temporal data out_channels = number of convolutional filters to use cnnfilter_size = kernel size for CNN cnnfilter_stride = stride of kernel cnn_padding = padding to use for data use_bias = whether CNN layer has a bias term use_maxpool - whether to use MAX Pooling pool_size = MAX pooling kernel size pool_stride = MAX pooling stride for kernel use_batchnorm = use batchnorm if True eps = batchnorm epsion momentum = batchnorm momentum affine = batchnorm affine (True or False) dropout = if nonzero, probability of neuron dropping out """ super(EEG_CNN_BuildingBlock, self).__init__() #useful params to keep self.use_batchnorm = use_batchnorm self.dropout = dropout self.CNN_bias = use_bias self.use_maxpool = use_maxpool #CNN Layer if cnn_pad == 'same': #for keeping dimensionality of input/output the same cnn_pad = int(((cnnfilter_stride - 1) * seq_len - cnnfilter_stride + cnnfilter_size) // 2) self.CNN = nn.Conv1d(in_channels, out_channels, cnnfilter_size, stride=cnnfilter_stride, padding=cnn_pad, bias=use_bias, padding_mode='zeros') #for use in chaining to another model self.Lout = int( (seq_len + 2 * self.CNN.padding[0] - self.CNN.dilation[0] * (self.CNN.kernel_size[0] - 1) - 1) / self.CNN.stride[0] + 1) #initialize the CNN weights with Xavier Norm Init nn.init.xavier_normal_(self.CNN.weight.data) if self.CNN_bias: self.CNN.bias.data.uniform_(0, 1) #RELU Layer self.RELU = nn.ReLU() #Batchnorm Layer if use_batchnorm: self.BatchNorm = nn.BatchNorm1d(out_channels, eps, momentum, affine) #Dropout Layer if dropout > 0: self.DropOut = nn.Dropout(dropout) #MaxPool Layer if use_maxpool: self.MaxPool = nn.MaxPool1d(pool_size, pool_stride) if type(cnn_pad) is int: #for use in chaining to another model. self.Lout = int((self.Lout + 2 * self.MaxPool.padding - self.MaxPool.dilation * (self.MaxPool.kernel_size - 1) - 1) / self.MaxPool.stride + 1) else: self.Lout = int(seq_len) self.out_channels = self.CNN.out_channels
def __init__(self, inplace=False) -> None: super().__init__() self.maxpool = nn.MaxPool1d(3) self.relu = nn.ReLU(inplace=inplace)
def test_maxpool(self): x = Variable(torch.randn(20, 16, 50)) self.assertONNXExpected(export_to_string(nn.MaxPool1d(3, stride=2), x))
u.Flatten(), nn.Linear(256, 1024), nn.ReLU(), nn.Dropout(p = 0.5), nn.Linear(1024, 1024), nn.ReLU(), nn.Dropout(p = 0.5), nn.Linear(1024, 208), nn.Softmax(), ), 'model.0' : nn.Sequential( nn.Conv1d(227, padding = (0,1), kernel_size = (1,3), stride = 1, out_channels = 100), nn.ReLU(), nn.MaxPool1d(padding = (0,0), kernel_size = (1,3), stride = 2), nn.Conv1d(100, padding = (0,1), kernel_size = (1,3), out_channels = 64, stride = 1), nn.ReLU(), nn.MaxPool1d(padding = (0,0), kernel_size = (1,2), stride = 2), nn.Conv1d(64, padding = (0,1), kernel_size = (1,3), out_channels = 128, stride = 1), nn.ReLU(), nn.MaxPool1d(padding = (0,0), kernel_size = (1,2), stride = 2), nn.Conv1d(128, padding = (0,1), kernel_size = (1,3), out_channels = 128, stride = 1), nn.ReLU(), nn.MaxPool1d(padding = (0,0), kernel_size = (1,2), stride = 2), nn.Conv1d(128, padding = (0,1), kernel_size = (1,3), out_channels = 128, stride = 1), nn.ReLU(),
def __init__(self, dim, inplanes, planes, downsample, use_gn, lr_mult, use_out, out_bn, sync_bn, whiten_type, temperature, with_gc, value_split, gc_beta): assert dim in [1, 2, 3], "dim {} is not supported yet".format(dim) #assert whiten_type in ['channel', 'spatial'] if dim == 3: conv_nd = nn.Conv3d if downsample: max_pool = nn.MaxPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2)) else: max_pool = None bn_nd = nn.BatchNorm3d elif dim == 2: conv_nd = nn.Conv2d if downsample: max_pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)) else: max_pool = None bn_nd = nn.BatchNorm2d else: conv_nd = nn.Conv1d if downsample: max_pool = nn.MaxPool1d(kernel_size=2, stride=2) else: max_pool = None bn_nd = nn.BatchNorm1d super(_NonLocalNd_bn, self).__init__() self.conv_query = conv_nd(inplanes, planes, kernel_size=1) self.conv_key = conv_nd(inplanes, planes, kernel_size=1) if use_out: if with_gc and value_split: self.conv_value = conv_nd(inplanes, planes, kernel_size=1) self.conv_gc_value = conv_nd(inplanes, planes, kernel_size=1) self.conv_out = conv_nd(planes, inplanes, kernel_size=1, bias=False) else: self.conv_value = conv_nd(inplanes, planes, kernel_size=1) self.conv_out = conv_nd(planes, inplanes, kernel_size=1, bias=False) else: if with_gc and value_split: self.conv_value = conv_nd(inplanes, inplanes, kernel_size=1, bias=False) self.conv_gc_value = conv_nd(inplanes, inplanes, kernel_size=1, bias=False) self.conv_out = None else: self.conv_value = conv_nd(inplanes, inplanes, kernel_size=1, bias=False) self.conv_out = None if out_bn: if sync_bn: self.out_bn = SyncNorm2d(inplanes) else: self.out_bn = nn.BatchNorm2d(inplanes) else: self.out_bn = None if with_gc: self.conv_mask = conv_nd(inplanes, 1, kernel_size=1) if 'bn_affine' in whiten_type: self.key_bn_affine = nn.BatchNorm1d(planes) self.query_bn_affine = nn.BatchNorm1d(planes) if 'bn' in whiten_type: self.key_bn = nn.BatchNorm1d(planes, affine=False) self.query_bn = nn.BatchNorm1d(planes, affine=False) self.softmax = nn.Softmax(dim=2) self.downsample = max_pool # self.norm = nn.GroupNorm(num_groups=32, num_channels=inplanes) if use_gn else InPlaceABNSync(num_features=inplanes) self.gamma = nn.Parameter(torch.zeros(1)) if gc_beta: self.beta = nn.Parameter(torch.zeros(1)) self.scale = math.sqrt(planes) self.whiten_type = whiten_type self.temperature = temperature self.with_gc = with_gc self.value_split = value_split self.gc_beta = gc_beta self.reset_parameters() self.reset_lr_mult(lr_mult)
def forward(self, char_embeddings): x_conv = self.conv1d(char_embeddings) x_conv_out = nn.MaxPool1d(nn.ReLU(x_conv)) return x_conv_out
def __init__(self, text_shape, target_shape, num_classes, embedding_vector, nonlin=nn.ReLU): super(BiLSTM, self).__init__() _, self.seq_len, self.embedding_dim = text_shape _, self.tar_len, _ = target_shape text_vector, target_vector = embedding_vector self.hidden_size = 32 self.num_layers = 1 # input: (m, seq_len), (m, tar_len) self.embedding_text = nn.Sequential( OrderedDict([ ('embed1a', nn.Embedding.from_pretrained(text_vector, freeze=False)), ])) self.embedding_target = nn.Sequential( OrderedDict([ ('embed1b', nn.Embedding.from_pretrained(target_vector, freeze=False)), ])) # output: (m, seq_len, embedding_dim), (m, tar_len, embedding_dim) # input: (m, seq_len, embedding_dim), (m, tar_len, embedding_dim) self.bilstm = nn.Sequential( OrderedDict([ ('bilstm2', nn.LSTM(self.embedding_dim, self.hidden_size, num_layers=self.num_layers, bidirectional=True, batch_first=True)), ])) # output: (m, seq_len, hidden_size * 2), (m, tar_len, hidden_size * 2) # input: (m, hidden_size * 2, seq_len), (m, hidden_size * 2, tar_len) self.pooling_text = nn.Sequential( OrderedDict([ ('maxpool3a', nn.MaxPool1d(self.seq_len, stride=1)), ])) self.pooling_target = nn.Sequential( OrderedDict([ ('maxpool3b', nn.MaxPool1d(self.tar_len, stride=1)), ])) # output: (m, hidden_size * 2, 1), (m, hidden_size * 2, 1) # input: (m, hidden_size * 2, 1), (m, hidden_size * 2, 1) self.nonlin = nn.Sequential(OrderedDict([ ('nonlin4', nonlin()), ])) # output: (m, hidden_size * 2, 1), (m, hidden_size * 2, 1) # input: (m, hidden_size * 4) self.classifier = nn.Sequential( OrderedDict([ ('dropout5', nn.Dropout(0.5)), ('fc5', nn.Linear(self.hidden_size * 4, num_classes)), ]))
def __init__(self, vocap_size, vector_size): super(Embed, self).__init__() self.embedding = nn.Embedding(vocab_size, vector_size) self.conv1 = nn.Conv1d(vector_size, 32, 8) self.maxpool = nn.MaxPool1d(2) self.relu = nn.ReLU(inplace=True)
def __init__(self, input_shape, out_shape, kernel_size, channel_size): super(KernelBigVGGDK, self).__init__() self.out_shape = out_shape self.input_shape = input_shape self.kernel_size = kernel_size self.padding = int(self.kernel_size / 2) self.max_pool = 2 self.channel_size = channel_size self.conv1_channels = self.channel_size self.conv2_channels = self.conv1_channels * 2 self.conv3_channels = self.conv2_channels * 2 self.conv4_channels = self.conv3_channels * 2 self.conv5_channels = self.conv4_channels * 2 num_features = input_shape ########## # STEP 1 # ########## self.bn0 = nn.BatchNorm1d(num_features=1).to(device) # Convolutions + BN + MP self.conv1 = nn.Conv1d(1, self.conv1_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.mp1 = nn.MaxPool1d(self.max_pool).to(device) num_features = int(num_features / 2) self.bn1 = nn.BatchNorm1d(num_features=self.conv1_channels).to(device) ########## # STEP 2 # ########## self.conv2_1 = nn.Conv1d(self.conv1_channels, self.conv2_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.conv2_2 = nn.Conv1d(self.conv2_channels, self.conv3_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.mp2 = nn.MaxPool1d(self.max_pool).to(device) num_features = int(num_features / 2) self.bn2 = nn.BatchNorm1d(num_features=self.conv3_channels).to(device) ########## # STEP 3 # ########## self.conv3_1 = nn.Conv1d(self.conv3_channels, self.conv4_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.conv3_2 = nn.Conv1d(self.conv4_channels, self.conv5_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.mp3 = nn.MaxPool1d(self.max_pool).to(device) num_features = int(num_features / 2) self.bn3 = nn.BatchNorm1d(num_features=self.conv5_channels).to(device) ########## # STEP 4 # ########## self.conv4_1 = nn.Conv1d(self.conv5_channels, self.conv5_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.conv4_2 = nn.Conv1d(self.conv5_channels, self.conv5_channels, kernel_size=self.kernel_size, padding=self.padding).to(device) num_features = num_features + 2 * self.padding - 1 * ( self.kernel_size - 1) self.mp4 = nn.MaxPool1d(self.max_pool).to(device) num_features = int(num_features / 2) self.bn4 = nn.BatchNorm1d(num_features=self.conv5_channels).to(device) ########### # Dropout # ########### self.drop_out = nn.Dropout(p=0.5) self.drop_out2 = nn.Dropout(p=0.5) ################### # FULLY CONNECTED # ################### self.fc4 = torch.nn.Linear( int(self.conv5_channels * num_features) + 256, 256).to(device) self.fc5 = torch.nn.Linear(256, self.out_shape).to(device)
def MaxPool1d(kernel_size=2): m = nn.MaxPool1d(kernel_size, padding=1) return m
def __init__(self, hidden_size, projection_size=128, K=16, num_gru_layers=2, max_pool_kernel_size=2, is_post=False): super(CBHG, self).__init__() self.hidden_size = hidden_size self.num_gru_layers = num_gru_layers self.projection_size = projection_size self.convbank_list = nn.ModuleList() self.convbank_list.append( nn.Conv1d(in_channels=projection_size, out_channels=hidden_size, kernel_size=1, padding=int(np.floor(1 / 2)))) for i in range(2, K + 1): self.convbank_list.append( nn.Conv1d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=i, padding=int(np.floor(i / 2)))) self.batchnorm_list = nn.ModuleList() for i in range(1, K + 1): self.batchnorm_list.append(nn.BatchNorm1d(hidden_size)) convbank_outdim = hidden_size * K if is_post: self.conv_projection_1 = nn.Conv1d(in_channels=convbank_outdim, out_channels=hidden_size * 2, kernel_size=3, padding=int(np.floor(3 / 2))) self.conv_projection_2 = nn.Conv1d(in_channels=hidden_size * 2, out_channels=projection_size, kernel_size=3, padding=int(np.floor(3 / 2))) self.batchnorm_proj_1 = nn.BatchNorm1d(hidden_size * 2) else: self.conv_projection_1 = nn.Conv1d(in_channels=convbank_outdim, out_channels=hidden_size, kernel_size=3, padding=int(np.floor(3 / 2))) self.conv_projection_2 = nn.Conv1d(in_channels=hidden_size, out_channels=projection_size, kernel_size=3, padding=int(np.floor(3 / 2))) self.batchnorm_proj_1 = nn.BatchNorm1d(hidden_size) self.batchnorm_proj_2 = nn.BatchNorm1d(projection_size) self.max_pool = nn.MaxPool1d(max_pool_kernel_size, stride=1, padding=1) self.highway = Highwaynet(self.projection_size) self.gru = nn.GRU(self.projection_size, self.hidden_size, num_layers=2, batch_first=True, bidirectional=True)
def __init__(self, maxpool): super(Net, self).__init__() self.down = nn.MaxPool1d(maxpool, stride=maxpool)
def __init__(self, input_channels, output_channels, m_word,kernel_size=5): super(CNN, self).__init__() self.conv = nn.Conv1d(input_channels, output_channels, kernel_size) self.maxp = nn.MaxPool1d(m_word-kernel_size+1)
def __init__( self, block, layers, in_channels=4, num_classes=2, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None, freeze_schedule=None, freeze_milestones=None, pool_size=1, ): super(ResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm1d self._norm_layer = norm_layer self.inplanes = 64 self.dilation = 1 if replace_stride_with_dilation is None: # each element in the tuple indicates if we should replace # the 2x2 stride with a dilated convolution instead replace_stride_with_dilation = [False, False, False] if len(replace_stride_with_dilation) != 3: raise ValueError( "replace_stride_with_dilation should be None " "or a 3-element tuple, got {}".format(replace_stride_with_dilation) ) self.groups = groups self.base_width = width_per_group self.conv1 = nn.Conv1d( in_channels, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False ) self.bn1 = norm_layer(self.inplanes) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1) self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer( block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0] ) self.layer3 = self._make_layer( block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1] ) self.layer4 = self._make_layer( block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2] ) self.avgpool = nn.AdaptiveAvgPool1d((pool_size)) self.fc = nn.Linear(512 * pool_size * block.expansion, num_classes) self.freeze_schedule = freeze_schedule self.freeze_milestones = freeze_milestones for m in self.modules(): if isinstance(m, nn.Conv1d): nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") elif isinstance(m, (nn.BatchNorm1d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0)
def __init__(self, framenum=120): super(VAMetric2, self).__init__() self.mp = nn.MaxPool1d(framenum) self.vfc = nn.Linear(1024, 128) self.fc = nn.Linear(128, 96) self.init_params()
def test_maxpool(self): x = torch.randn(20, 16, 50) self.assertONNX(nn.MaxPool1d(3, stride=2), x)
def __init__(self, input_dim=40, hidden_dim=256, speaker_num=1941, encoder='pBiLSTM', decoder='linear', pooling='max'): super().__init__() # self.cnn0 = nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=3, stride=1, padding=1) # self.cnn1 = nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=3, stride=1, padding=1) # self.cnn2 = nn.Conv1d(in_channels=input_dim, out_channels=input_dim, kernel_size=3, stride=1, padding=1) # self.ReLU = nn.LeakyReLU(0.01) ## Encoder self.encoder = encoder self.decoder = decoder if self.encoder == 'pBiLSTM' or self.encoder == 'BiLSTM': self.batch_first = False else: self.batch_first = True if encoder == 'pBiLSTM': self.seq_pool = SequencePooling() self.rnns_encoder = nn.ModuleList([ LSTM_State0(input_size=input_dim, hidden_size=hidden_dim, batch_first=False, bidirectional=True), LSTM_State0(input_size=hidden_dim * 4, hidden_size=hidden_dim, batch_first=False, bidirectional=True), LSTM_State0(input_size=hidden_dim * 4, hidden_size=hidden_dim, batch_first=False, bidirectional=True) ]) self.final_rnn_encode = LSTM_State0(input_size=hidden_dim * 4, hidden_size=hidden_dim, batch_first=False, bidirectional=True) # self.keys = nn.Linear(in_features=hidden_dim*2, # out_features=key_dim) # self.values = nn.Linear(in_features=hidden_dim*2, # out_features=key_dim) # Output is (L,B,D) (length in time, batch, feature dimension) (D=hidden_dim*2) elif encoder == 'BiLSTM': self.rnns_encoder = nn.ModuleList([ LSTM_State0(input_size=input_dim, hidden_size=hidden_dim, batch_first=False, bidirectional=True), LSTM_State0(input_size=hidden_dim * 2, hidden_size=hidden_dim, batch_first=False, bidirectional=True), LSTM_State0(input_size=hidden_dim * 2, hidden_size=hidden_dim, batch_first=False, bidirectional=True) ]) elif encoder == 'CNN2D': self.conv1 = nn.Conv2d(40, 96, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1) self.conv3 = nn.Conv2d(96, 96, kernel_size=3, stride=1, padding=1) self.conv4 = nn.Conv2d(96, 192, kernel_size=3, stride=1, padding=1) self.conv5 = nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1) self.conv6 = nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1) self.conv7 = nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1) self.conv8 = nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1) self.conv9 = nn.Conv2d(192, 192, kernel_size=3, stride=1, padding=1) self.conv10 = nn.Conv2d(192, 46, kernel_size=3, stride=1, padding=1) # self.conv_encoder = nn.ModuleList([ # nn.Conv2d(in_channels=input_dim, out_channels=32, kernel_size=4, stride=2, padding=2), # nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2, padding=2)]) self.maxpool = nn.MaxPool2d(4, stride=2, padding=2) print('Not Implemented') elif encoder == 'CNN1D': self.conv1 = nn.Conv1d(40, 96, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv1d(96, 96, kernel_size=3, stride=1, padding=1) self.conv3 = nn.Conv1d(96, 96, kernel_size=3, stride=1, padding=1) self.conv4 = nn.Conv1d(96, 192, kernel_size=3, stride=1, padding=1) self.conv5 = nn.Conv1d(192, 192, kernel_size=3, stride=1, padding=1) self.conv6 = nn.Conv1d(192, 192, kernel_size=3, stride=1, padding=1) self.conv7 = nn.Conv1d(192, hidden_dim * 2, kernel_size=3, stride=1, padding=1) self.conv8 = nn.Conv1d(hidden_dim * 2, hidden_dim * 2, kernel_size=3, stride=1, padding=1) self.conv9 = nn.Conv1d(hidden_dim * 2, hidden_dim * 2, kernel_size=3, stride=1, padding=1) self.conv10 = nn.Conv1d(hidden_dim * 2, hidden_dim * 2, kernel_size=3, stride=1, padding=1) self.leaky_ReLU = nn.LeakyReLU(0.2) self.maxpool = nn.MaxPool1d(4, stride=2, padding=2) # print('Not Implemented') ## Decoder if decoder == 'linear': self.linear = nn.Linear(in_features=hidden_dim * 2, out_features=speaker_num) self.pooling = cust_Pool(pooling, self.batch_first) elif decoder == 'MLP': self.mlp = nn.ModuleList([ nn.Linear(in_features=hidden_dim * 2, out_features=hidden_dim * 2), nn.Linear(in_features=hidden_dim * 2, out_features=hidden_dim * 2), nn.Linear(in_features=hidden_dim * 2, out_features=hidden_dim) ]) self.mlp_fin = nn.Linear(in_features=hidden_dim, out_features=speaker_num) self.pooling = cust_Pool(pooling, self.batch_first) self.activation = nn.LeakyReLU(0.2) elif decoder == 'FC': # Since input is variable length, can this be done? # self.resizer = print('Not Implemented') elif decoder == 'CNN2D': self.conv_decoder = nn.ModuleList([ nn.Conv2d(in_channels=1, out_channels=hidden_dim * 2, kernel_size=4, stride=2, padding=2), nn.Conv2d(in_channels=hidden_dim * 2, out_channels=hidden_dim * 2, kernel_size=4, stride=2, padding=2), nn.Conv2d(in_channels=hidden_dim * 2, out_channels=hidden_dim, kernel_size=4, stride=2, padding=2) ]) self.relu = nn.ReLU() self.linear = nn.Linear(in_features=hidden_dim, out_features=speaker_num) if pooling == 'max': self.pool = nn.MaxPool2d(4, stride=2) elif pooling == 'mean': self.pool = nn.AvgPool2d(4, stride=2) self.pooling = cust_Pool(pooling, self.batch_first) print('Not Implemented') elif decoder == 'CNN1D': self.conv_decoder = nn.ModuleList([ nn.Conv1d(in_channels=hidden_dim * 2, out_channels=hidden_dim * 2, kernel_size=4, stride=2, padding=2), nn.Conv1d(in_channels=hidden_dim * 2, out_channels=hidden_dim * 2, kernel_size=4, stride=2, padding=2), nn.Conv1d(in_channels=hidden_dim * 2, out_channels=hidden_dim, kernel_size=4, stride=2, padding=2) ]) self.relu = nn.ReLU() self.linear = nn.Linear(in_features=hidden_dim, out_features=speaker_num) if pooling == 'max': self.pool = nn.MaxPool1d(4, stride=2) elif pooling == 'mean': self.pool = nn.AvgPool1d(4, stride=2) self.pooling = cust_Pool(pooling, self.batch_first) print('Not Implemented') elif decoder == 'attention': self.attention = nn.Linear(in_features=hidden_dim * 2, out_features=speaker_num) self.weight_W = nn.Linear(in_features=hidden_dim * 2, out_features=1) self.softmax2 = nn.Softmax(dim=2) self.softmax1 = nn.Softmax(dim=1)
def test_maxpool_indices(self): x = torch.randn(20, 16, 50) self.assertONNX(nn.MaxPool1d(3, stride=2, return_indices=True), x)
def __init__(self, dataset, gconv=RGCNConv, latent_dim=[32, 32, 32, 32], num_relations=5, num_bases=2, regression=False, adj_dropout=0.2, force_undirected=False, side_features=False, n_side_features=0, multiply_by=1, use_graphnorm=False, model_type='IGMC', gconv_type='GCNConv'): if gconv_type == 'SAGEConv': super(IGMC, self).__init__(dataset, SAGEConv, latent_dim, regression, adj_dropout, force_undirected) # Default: GCNConv else: super(IGMC, self).__init__(dataset, GCNConv, latent_dim, regression, adj_dropout, force_undirected) self.multiply_by = multiply_by #convolutions self.convs = torch.nn.ModuleList() self.convs.append( gconv(dataset.num_features, latent_dim[0], num_relations, num_bases)) for i in range(0, len(latent_dim) - 1): self.convs.append( gconv(latent_dim[i], latent_dim[i + 1], num_relations, num_bases)) #norms self.use_graphnorm = use_graphnorm self.model_type = model_type if use_graphnorm: self.norms = torch.nn.ModuleList() for i in range(len(latent_dim)): self.norms.append(GraphNorm(latent_dim[i])) if model_type == 'MaxPoolIGMC': self.maxpool1d = nn.MaxPool1d(4) self.lin1 = Linear(int(2 * sum(latent_dim) / 4), 128) else: self.lin1 = Linear(2 * sum(latent_dim), 128) self.side_features = side_features if side_features: self.lin1 = Linear(2 * sum(latent_dim) + n_side_features, 128) #LSTM Attention params self.hidden_size = 16 self.bi_lstm = torch.nn.LSTM(input_size=32, hidden_size=self.hidden_size, bidirectional=True) self.lin3 = Linear(32, 1).cuda() self.softmax = torch.nn.Softmax(dim=0) self.num_layers = len(latent_dim)
def __init__(self, embedding, pep_length, tcr_length): super(Net, self).__init__() ## embedding layer self.num_amino = len(embedding) self.embedding_dim = len(embedding[0]) self.embedding = nn.Embedding(self.num_amino, self.embedding_dim, padding_idx=self.num_amino-1).\ from_pretrained(torch.FloatTensor(embedding), freeze = False) ## peptide encoding layer self.size_hidden1_cnn = SIZE_HIDDEN1_CNN self.size_hidden2_cnn = SIZE_HIDDEN2_CNN self.size_kernel1 = SIZE_KERNEL1 self.size_kernel2 = SIZE_KERNEL2 self.size_padding = (self.size_kernel1 - 1) / 2 self.encode_pep = nn.Sequential( nn.Dropout(0.3), nn.Conv1d(self.embedding_dim, self.size_hidden1_cnn, kernel_size=self.size_kernel1), nn.BatchNorm1d(self.size_hidden1_cnn), nn.ReLU(True), nn.MaxPool1d(kernel_size=self.size_kernel1, stride=1, padding=self.size_padding), nn.Conv1d(self.size_hidden1_cnn, self.size_hidden2_cnn, kernel_size=self.size_kernel2), nn.BatchNorm1d(self.size_hidden2_cnn), nn.ReLU(True), nn.MaxPool1d(kernel_size=self.size_kernel2)) ## trc encoding layer self.encode_tcr = nn.Sequential( nn.Dropout(0.3), nn.Conv1d(self.embedding_dim, self.size_hidden1_cnn, kernel_size=self.size_kernel1), nn.BatchNorm1d(self.size_hidden1_cnn), nn.ReLU(True), nn.MaxPool1d(kernel_size=self.size_kernel1, stride=1, padding=self.size_padding), nn.Conv1d(self.size_hidden1_cnn, self.size_hidden2_cnn, kernel_size=self.size_kernel2), nn.BatchNorm1d(self.size_hidden2_cnn), nn.ReLU(True), nn.MaxPool1d(kernel_size=self.size_kernel2)) ## dense layer at the end self.net_pep_dim = self.size_hidden2_cnn * ( (pep_length - self.size_kernel1 + 1 - self.size_kernel2 + 1) / self.size_kernel2) self.net_tcr_dim = self.size_hidden2_cnn * ( (tcr_length - self.size_kernel1 + 1 - self.size_kernel2 + 1) / self.size_kernel2) self.net = nn.Sequential( nn.Dropout(0.3), nn.Linear(self.net_pep_dim + self.net_tcr_dim, 32), nn.ReLU(), nn.Linear(32, 16), nn.ReLU(), nn.Linear(16, 2), nn.LogSoftmax(1))
def __init__(self, **kwargs): super(CNN_GRU, self).__init__() self.MODEL = kwargs["MODEL"] self.BATCH_SIZE = kwargs["BATCH_SIZE"] self.MAX_SENT_LEN = kwargs["MAX_SENT_LEN"] self.WORD_DIM = kwargs["WORD_DIM"] self.VOCAB_SIZE = kwargs["VOCAB_SIZE"] self.CLASS_SIZE = kwargs["CLASS_SIZE"] self.FILTERS = kwargs["FILTERS"] self.FILTER_NUM = kwargs["FILTER_NUM"] self.DROPOUT_PROB = kwargs["DROPOUT_PROB"] self.IN_CHANNEL = 1 self.USE_CUDA = True self.HIDDEN_DIM = kwargs["HIDDEN_DIM"] self.NUM_LAYERS = kwargs["NUM_LAYERS"] self.bidirectional = True #kwargs["BIDIRECTIONAL"] self.embedding = nn.Embedding(self.VOCAB_SIZE + 2, self.WORD_DIM, padding_idx=self.VOCAB_SIZE + 1) if self.MODEL == "static" or self.MODEL == "non-static" or self.MODEL == "multichannel": self.WV_MATRIX = kwargs["WV_MATRIX"] self.embedding.weight.data.copy_(torch.from_numpy(self.WV_MATRIX)) if self.MODEL == "static": self.embedding.weight.requires_grad = False elif self.MODEL == "multichannel": self.embedding2 == nn.Embedding(self.VOCAB_SIZE + 2, self_WORD_DIM, padding_idx=self.VOCAB_SIZE + 1) self.embedding2.weight.data.copy_(torch.from_numpy(self.WV_MATRIX)) self.embedding2.weight.requires_grad = False self.IN_CHANNEL = 2 conv_blocks = [] for filter_size in self.FILTERS: maxpool_kernel_size = self.MAX_SENT_LEN - filter_size + 1 if ConvMethod == "in_channel_is_embedding_dim": conv1d = nn.Conv1d(in_channels=self.WORD_DIM, out_channels=self.FILTER_NUM, kernel_size=filter_size, stride=1) conv2 = nn.Conv1d(in_channels=self.FILTER_NUM, out_channels=self.FILTER_NUM, kernel_size=filter_size / 2, stride=1) else: conv1d = nn.Conv1d(in_channels=1, out_channels=self.FILTER_NUM, kernel_size=filter_size * self.WORD_DIM, stride=self.WORD_DIM) conv2 = nn.Conv1d(in_channels=self.FILTER_NUM, out_channels=self.FILTER_NUM, kernel_size=filter_size * self.WORD_DIM, stride=self.WORD_DIM) component = nn.Sequential( conv1d, nn.ReLU(), nn.MaxPool1d(kernel_size=maxpool_kernel_size), ) if self.USE_CUDA: component = component.cuda() conv_blocks.append(component) self.conv_blocks = nn.ModuleList(conv_blocks) self.GRU = nn.GRU(self.FILTER_NUM * len(self.FILTERS), self.HIDDEN_DIM, dropout=self.DROPOUT_PROB, num_layers=self.NUM_LAYERS, bidirectional=True) self.hidden2label = nn.Linear(self.HIDDEN_DIM * 2, self.CLASS_SIZE) self.hidden = self.init_hidden()
def __init__(self): super(Conv1ResizePoold, self).__init__() self.pool = nn.MaxPool1d(3, stride=2, padding=1) self.conv = nn.Conv2d(2, 2, kernel_size=3, stride=1, padding=1)
def __init__(self, noise_dim=64, local_features_dim=64): super(GlobalPathway, self).__init__() channel_encoder = [64, 64, 128, 256, 512] channel_decoder_feat = [64, 32, 16, 8] channel_decoder = [512, 256, 128, 64] channel_decoder_conv = [64, 32] self.noise_dim = noise_dim ## encoder blocks # img size = 128x128 self.conv0 = Sequential( ConvBlock(in_channels=3, out_channels=channel_encoder[0], kernel_size=7, stride=1, padding=3, activation=nn.LeakyReLU(), use_batchnorm=True), ResidualBlock(in_channels=channel_encoder[0], kernel_size=7, padding=3, activation=nn.LeakyReLU())) # img size = 64x64 self.conv1 = Sequential( ConvBlock(in_channels=channel_encoder[0], out_channels=channel_encoder[1], kernel_size=5, stride=2, padding=2, activation=nn.LeakyReLU(), use_batchnorm=True), ResidualBlock(in_channels=channel_encoder[0], kernel_size=5, padding=2, activation=nn.LeakyReLU())) # img size = 32x32 self.conv2 = Sequential( ConvBlock(in_channels=channel_encoder[1], out_channels=channel_encoder[2], kernel_size=3, stride=2, padding=1, activation=nn.LeakyReLU(), use_batchnorm=True), ResidualBlock(in_channels=channel_encoder[2], kernel_size=3, padding=1, activation=nn.LeakyReLU())) # img size = 16x16 self.conv3 = Sequential( ConvBlock(in_channels=channel_encoder[2], out_channels=channel_encoder[3], kernel_size=3, stride=2, padding=1, activation=nn.LeakyReLU(), use_batchnorm=True), ResidualBlock(in_channels=channel_encoder[3], kernel_size=3, padding=1, activation=nn.LeakyReLU())) # img size = 8x8 self.conv4 = Sequential( ConvBlock(in_channels=channel_encoder[3], out_channels=channel_encoder[4], kernel_size=3, stride=2, padding=1, activation=nn.LeakyReLU(), use_batchnorm=True), *[ ResidualBlock(in_channels=channel_encoder[4], kernel_size=3, padding=1, activation=nn.LeakyReLU()) for i in range(4) ]) self.fc1 = nn.Linear(channel_encoder[4] * 8 * 8, 512) self.fc2 = nn.MaxPool1d(2, 2, 0) ## decoder blocks # 8x8 self.feat8 = DeConvBlock(in_channels=channel_encoder[4] // 2 + self.noise_dim, out_channels=channel_decoder_feat[0], kernel_size=8, stride=1, padding=0, output_padding=0, activation=nn.ReLU(), use_batchnorm=True) self.decode8 = ResidualBlock(in_channels=self.feat8.out_channels + self.conv4.out_channels, kernel_size=3, stride=1, padding=1, activation=nn.LeakyReLU()) self.reconstruct8 = Sequential(*[ ResidualBlock(in_channels=self.feat8.out_channels + self.conv4.out_channels, kernel_size=3, stride=1, padding=1, activation=nn.LeakyReLU()) for i in range(2) ]) # 16x16 self.deconv16 = DeConvBlock(in_channels=self.reconstruct8.out_channels, out_channels=channel_decoder[0], kernel_size=3, stride=2, padding=1, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.decode16 = ResidualBlock(in_channels=self.conv3.out_channels, activation=nn.LeakyReLU()) self.reconstruct16 = Sequential(*[ ResidualBlock(in_channels=self.deconv16.out_channels + self.decode16.out_channels, activation=nn.LeakyReLU()) for i in range(2) ]) # 32x32 self.feat32 = DeConvBlock(in_channels=channel_decoder_feat[0], out_channels=channel_decoder_feat[1], kernel_size=3, stride=4, padding=0, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.deconv32 = DeConvBlock( in_channels=self.reconstruct16.out_channels, out_channels=channel_decoder[1], kernel_size=3, stride=2, padding=1, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.decode32 = ResidualBlock(in_channels=self.feat32.out_channels + self.conv2.out_channels + 3, activation=nn.LeakyReLU()) self.reconstruct32 = Sequential(*[ ResidualBlock(in_channels=self.feat32.out_channels + self.conv2.out_channels + 3 + self.deconv32.out_channels, activation=nn.LeakyReLU()) for i in range(2) ]) self.output32 = ConvBlock(in_channels=self.reconstruct32.out_channels, out_channels=3, kernel_size=3, stride=1, padding=1, activation=nn.Tanh(), init_weight=False) # 64x64 self.feat64 = DeConvBlock(in_channels=channel_decoder_feat[1], out_channels=channel_decoder_feat[2], kernel_size=3, stride=2, padding=1, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.deconv64 = DeConvBlock( in_channels=self.reconstruct32.out_channels, out_channels=channel_decoder[2], kernel_size=3, stride=2, padding=1, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.decode64 = ResidualBlock(in_channels=self.feat64.out_channels + self.conv1.out_channels + 3, kernel_size=5, activation=nn.LeakyReLU()) self.reconstruct64 = Sequential(*[ ResidualBlock(in_channels=self.feat64.out_channels + self.conv1.out_channels + 3 + self.deconv64.out_channels + 3, activation=nn.LeakyReLU()) for i in range(2) ]) self.output64 = ConvBlock(in_channels=self.reconstruct64.out_channels, out_channels=3, kernel_size=3, stride=1, padding=1, activation=nn.Tanh(), init_weight=False) # 128x128 self.feat128 = DeConvBlock(in_channels=channel_decoder_feat[2], out_channels=channel_decoder_feat[3], kernel_size=3, stride=2, padding=1, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.deconv128 = DeConvBlock( in_channels=self.reconstruct64.out_channels, out_channels=channel_decoder[3], kernel_size=3, stride=2, padding=1, output_padding=1, activation=nn.ReLU(), use_batchnorm=True) self.decode128 = ResidualBlock(in_channels=self.feat128.out_channels + self.conv0.out_channels + 3, kernel_size=7, activation=nn.LeakyReLU()) self.reconstruct128 = ResidualBlock( in_channels=self.feat128.out_channels + self.conv0.out_channels + 3 + self.deconv128.out_channels + 3 + local_features_dim + 3, kernel_size=5, activation=nn.LeakyReLU()) self.decode_conv0 = Sequential( ConvBlock(in_channels=self.reconstruct128.out_channels, out_channels=channel_decoder_conv[0], kernel_size=5, stride=1, padding=2, activation=nn.LeakyReLU(), use_batchnorm=True), ResidualBlock(channel_decoder_conv[0], kernel_size=3)) self.decode_conv1 = ConvBlock(in_channels=channel_decoder_conv[0], out_channels=channel_decoder_conv[1], kernel_size=3, stride=1, padding=1, activation=nn.LeakyReLU(), use_batchnorm=True) self.output128 = ConvBlock(in_channels=channel_decoder_conv[1], out_channels=3, kernel_size=3, stride=1, padding=1, activation=nn.Tanh(), init_weight=False)