def __init__(self, in_channels, inter_channels=None, dimension=2, sub_sample=True, bn_layer=True, is_training=True): assert dimension in [1, 2, 3], "Unknown dimension..." self.in_channels = in_channels self.inter_channels = inter_channels self.sub_sample = sub_sample if self.inter_channels is None: self.inter_channels = in_channels // 2 if self.inter_channels == 0: self.inter_channels = 1 self.g = ConvBlock(self.inter_channels, dimension, 1) self.W = ConvBlock(self.in_channels, dimension, 1, use_bn=bn_layer, is_training=is_training) if sub_sample: self.g_pool = MaxPool(dimension) self.phi = MaxPool(dimension)
def __init__(self, config): super(Discriminator, self).__init__() self.is_cuda = torch.cuda.is_available() N = int(config.nfc) self.head = ConvBlock(config.img_channel, N, config.kernel_size, 1, config.pad) self.body = nn.Sequential() for i in range(config.num_layers - 2): N = int(config.nfc / pow(2, (i + 1))) block = ConvBlock(max(2 * N, config.min_nfc), max(N, config.min_nfc), config.kernel_size, 1, config.pad) self.body.add_module('block%d' % (i + 1), block) # WGAN-GP discriminator has no activation at last layer self.tail = nn.Conv2d(max(N, config.min_nfc), 1, kernel_size=config.kernel_size, stride=1, padding=config.pad)
def __init__(self, config, num_heads, real): super(ACMDiscriminator, self).__init__() self.ans = real self.config = config self.is_cuda = torch.cuda.is_available() N = int(config.nfc) self.head = ConvBlock(config.img_channel, N, config.kernel_size, 1, config.pad) self.body = nn.Sequential() for i in range(config.num_layers - 2): N = int(config.nfc / pow(2, (i + 1))) block = ConvBlock(max(2 * N, config.min_nfc), max(N, config.min_nfc), config.kernel_size, 1, config.pad) self.body.add_module('block%d' % (i + 1), block) self.acm = CustomACM(num_heads=num_heads, num_features=max(N, config.min_nfc), orthogonal_loss=self.config.use_acm_oth) # WGAN-GP discriminator has no activation at last layer self.tail = nn.Conv2d(max(N, config.min_nfc), 1, kernel_size=config.kernel_size, stride=1, padding=config.pad)
class DRNSeg(object): def __init__(self, model_name, num_classes, use_non_local=True, use_bilinear=True, is_training=True): self.num_classes = num_classes self.use_bilinear = use_bilinear self.use_non_local = use_non_local if model_name == 'drn_c_26': self.base = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1], num_classes=num_classes, out_map=True, out_middle=True, is_training=is_training) else: raise Exception("Unknown architecture...") if use_non_local: self.nonLocal = NonLocalBlock2D(512) self.seg = ConvBlock(num_classes, kernel_size=1) # self.decoder = DecoderMudule(48, num_classes, use_bilinear=use_bilinear, is_training=is_training) def forward(self, inputs): inputs_shape = inputs.get_shape().as_list() height, width = inputs_shape[1], inputs_shape[2] _, end_points = self.base.forward(inputs) out = end_points[-1] encoder_features = end_points[1:3] if self.use_non_local: out = self.nonLocal(out) out = self.seg.forward(out) # out = self.decoder.forward(out, encoder_features) if self.use_bilinear: out = tf.image.resize_bilinear(out, (height, width)) else: out = tf.layers.conv2d_transpose( out, self.num_classes, 2, 2, padding='same', use_bias=False, kernel_initializer=tf.truncated_normal_initializer( stddev=0.02)) pred = tf.argmax(out, axis=3, name='prediction') pred = tf.expand_dims(pred, axis=3) return out, pred def __call__(self, inputs): return self.forward(inputs)
class BasicBlock(object): expansion = 1 def __init__(self, out_channels, stride=1, dilation=(1, 1), residual=True, is_training=True): self.out_channels = out_channels self.stride = stride self.residual = residual self.is_training = is_training self.conv_block1 = ConvBlock(out_channels, stride=stride, dilation=dilation[0], use_bn=True, use_act=True, is_training=is_training) self.conv_block2 = ConvBlock(out_channels, dilation=dilation[1], use_bn=True, is_training=is_training) def forward(self, inputs): in_channels = tf.shape(inputs)[-1] res = inputs out = self.conv_block1.forward(inputs) out = self.conv_block2.forward(out) if self.stride != 1 or in_channels != self.out_channels: downsample = ConvBlock(self.out_channels, kernel_size=1, stride=self.stride, is_training=self.is_training) res = downsample.forward(inputs) if self.residual: out += res out = tf.nn.relu(out) return out
def forward(self, inputs): in_channels = tf.shape(inputs)[-1] res = inputs out = self.conv_block1.forward(inputs) out = self.conv_block2.forward(out) if self.stride != 1 or in_channels != self.out_channels: downsample = ConvBlock(self.out_channels, kernel_size=1, stride=self.stride, is_training=self.is_training) res = downsample.forward(inputs) if self.residual: out += res out = tf.nn.relu(out) return out
def __init__(self, out_channels, stride=1, dilation=(1, 1), residual=True, is_training=True): self.out_channels = out_channels self.stride = stride self.residual = residual self.is_training = is_training self.conv_block1 = ConvBlock(out_channels, stride=stride, dilation=dilation[0], use_bn=True, use_act=True, is_training=is_training) self.conv_block2 = ConvBlock(out_channels, dilation=dilation[1], use_bn=True, is_training=is_training)
def __init__(self, config): super(Generator, self).__init__() self.is_cuda = torch.cuda.is_available() N = config.nfc self.head = ConvBlock(config.img_channel, N, config.kernel_size, 1, config.pad) self.body = nn.Sequential() for i in range(config.num_layers - 2): N = int(config.nfc / pow(2, (i + 1))) block = ConvBlock(max(2 * N, config.min_nfc), max(N, config.min_nfc), config.kernel_size, 1, config.pad) self.body.add_module('block%d' % (i + 1), block) self.tail = nn.Sequential( nn.Conv2d(max(N, config.min_nfc), config.img_channel, kernel_size=config.kernel_size, stride=1, padding=config.pad), nn.Tanh())
def __init__(self, model_name, num_classes, use_non_local=True, use_bilinear=True, is_training=True): self.num_classes = num_classes self.use_bilinear = use_bilinear self.use_non_local = use_non_local if model_name == 'drn_c_26': self.base = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1], num_classes=num_classes, out_map=True, out_middle=True, is_training=is_training) else: raise Exception("Unknown architecture...") if use_non_local: self.nonLocal = NonLocalBlock2D(512) self.seg = ConvBlock(num_classes, kernel_size=1)
def _build_layers(self): self.conv1 = tf.make_template('conv1', ConvBlock(kernel_size=9, filters=16, data_format=self.data_format, padding=self.padding)) self.conv2 = tf.make_template('conv2', ConvBlock(kernel_size=7, filters=32, data_format=self.data_format, padding=self.padding)) self.conv3 = tf.make_template('conv3', ConvBlock(kernel_size=5, filters=64, data_format=self.data_format, padding=self.padding)) self.conv4 = tf.make_template('conv4', ConvBlock(kernel_size=4, filters=64, data_format=self.data_format, padding=self.padding)) # if use_edges is true, then the output from the Generator must be 3 channels, where: # channel 0: probability of pixel is black i.e. background # channel 1: probability pixel is white i.e. the body of a cell # channel 2: 2, probability pixel is an edge of a cell # # if use_edges is false, then the output is a single channel where each value represents # the probability of a pixel is a body of a cell if self.use_edges: self.conv5 = tf.make_template('conv5', ConvBlock(kernel_size=1, filters=3, data_format=self.data_format, padding=self.padding, is_final=True, use_edges=True)) else: self.conv5 = tf.make_template('conv5', ConvBlock(kernel_size=1, filters=1, data_format=self.data_format, padding=self.padding, is_final=True))
def __init__(self, block, layers, num_classes=5, channels=(16, 32, 64, 128, 256, 512, 512, 512), out_map=False, out_middle=False, pool_size=32, is_training=True): self.out_map = out_map self.pool_size = pool_size self.out_middle = out_middle self.conv1 = ConvBlock(channels[0], kernel_size=7, use_bias=False, use_bn=True, use_act=True, is_training=is_training) self.layer1 = MakeLayer(BasicBlock, channels[0], layers[0], stride=1, is_training=is_training) self.layer2 = MakeLayer(BasicBlock, channels[1], layers[1], stride=2, is_training=is_training) self.layer3 = MakeLayer(block, channels[2], layers[2], stride=2, is_training=is_training) self.layer4 = MakeLayer(block, channels[3], layers[3], stride=2, is_training=is_training) self.layer5 = MakeLayer(block, channels[4], layers[4], dilation=2, new_level=False, is_training=is_training) if layers[5] == 0: self.layer6 = None else: self.layer6 = MakeLayer(block, channels[5], layers[5], dilation=4, new_level=False, is_training=is_training) if layers[6] == 0: self.layer7 = None else: self.layer7 = MakeLayer(BasicBlock, channels[6], layers[6], dilation=2, new_level=False, residual=False, is_training=is_training) if layers[7] == 0: self.layer8 = None else: self.layer8 = MakeLayer(BasicBlock, channels[7], layers[7], dilation=1, new_level=False, residual=False, is_training=is_training) if num_classes > 0: self.fc = ConvBlock(num_classes, kernel_size=1, is_training=is_training)
class DRN(object): def __init__(self, block, layers, num_classes=5, channels=(16, 32, 64, 128, 256, 512, 512, 512), out_map=False, out_middle=False, pool_size=32, is_training=True): self.out_map = out_map self.pool_size = pool_size self.out_middle = out_middle self.conv1 = ConvBlock(channels[0], kernel_size=7, use_bias=False, use_bn=True, use_act=True, is_training=is_training) self.layer1 = MakeLayer(BasicBlock, channels[0], layers[0], stride=1, is_training=is_training) self.layer2 = MakeLayer(BasicBlock, channels[1], layers[1], stride=2, is_training=is_training) self.layer3 = MakeLayer(block, channels[2], layers[2], stride=2, is_training=is_training) self.layer4 = MakeLayer(block, channels[3], layers[3], stride=2, is_training=is_training) self.layer5 = MakeLayer(block, channels[4], layers[4], dilation=2, new_level=False, is_training=is_training) if layers[5] == 0: self.layer6 = None else: self.layer6 = MakeLayer(block, channels[5], layers[5], dilation=4, new_level=False, is_training=is_training) if layers[6] == 0: self.layer7 = None else: self.layer7 = MakeLayer(BasicBlock, channels[6], layers[6], dilation=2, new_level=False, residual=False, is_training=is_training) if layers[7] == 0: self.layer8 = None else: self.layer8 = MakeLayer(BasicBlock, channels[7], layers[7], dilation=1, new_level=False, residual=False, is_training=is_training) if num_classes > 0: self.fc = ConvBlock(num_classes, kernel_size=1, is_training=is_training) def forward(self, inputs): end_points = [] out = self.conv1.forward(inputs) out = self.layer1.forward(out) end_points.append(out) out = self.layer2.forward(out) end_points.append(out) out = self.layer3.forward(out) end_points.append(out) out = self.layer4.forward(out) end_points.append(out) out = self.layer5.forward(out) end_points.append(out) if self.layer6 is not None: out = self.layer6.forward(out) end_points.append(out) if self.layer7 is not None: out = self.layer7.forward(out) end_points.append(out) if self.layer8 is not None: out = self.layer8.forward(out) end_points.append(out) if self.out_map: out = self.fc.forward(out) else: out = tf.layers.average_pooling2d(out, self.pool_size, self.pool_size) out = self.fc.forward(out) out = tf.reshape(out, (out.shape[0], -1)) if self.out_middle: return out, end_points else: return out