def __init__(self, cin, comp_graph): super(GepBlock, self).__init__() self.paths = len(comp_graph) self.relu = relu(True) for i in range(self.paths): setattr(self, 'path_%d' % i, GepNetLayer(cin, comp_graph[i])) self.convproj = conv2d(cin * self.paths, cin, ksize=1, use_relu=False)
def conv2d(cin, cout=None, ksize=3, stride=1, padding=None, dilation=None, groups=None, use_relu=True, use_bn=True, bn=NormType.Batch, bias=False): if cout is None: cout = cin if padding is None: padding = ksize // 2 if dilation is None: dilation = 1 if groups is None: groups = 1 layer = [ init_default( nn.Conv2d(cin, cout, ksize, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias), nn.init.kaiming_normal_) ] if use_bn: layer.append(batchnorm_2d(cout, norm_type=bn)) if use_relu: layer.append(relu(True)) return nn.Sequential(*layer)
def __init__(self, cin, comp_graph): super(Cell, self).__init__() self.n_branch = len(comp_graph) self.relu = relu(True) for i in range(self.n_branch): setattr(self, 'branch_%d' % i, Layer(cin, comp_graph[i])) self.convproj = conv2d(cin*self.n_branch, cin, ksize=1, use_relu=False)
def __init__(self, up_in_c:int, final_div:bool=True, blur:bool=False, leaky:float=None, self_attention:bool=False, **kwargs): super().__init__() self.shuf = PixelShuffle_ICNR(up_in_c, up_in_c//2, blur=blur, leaky=leaky, **kwargs) ni = up_in_c//2 nf = ni if final_div else ni//2 self.conv1 = conv_layer(ni, nf, leaky=leaky, **kwargs) self.conv2 = conv_layer(nf, nf, leaky=leaky, self_attention=self_attention, **kwargs) self.relu = relu(leaky=leaky)
def stem_blk(cin, cout=None, ksize=3, stride=1, use_relu=True, use_bn=True, bn=NormType.Batch, bias=False, pool='avg'): if cout is None: cout = cin padding = ksize // 2 layer = [ init_default( nn.Conv2d(cin, cout, ksize, stride=stride, padding=padding, bias=bias), nn.init.kaiming_normal_) ] if use_bn: layer.append(batchnorm_2d(cout, norm_type=bn)) if use_relu: layer.append(relu(True)) if pool == 'max': layer.append(nn.MaxPool2d(2, stride=2)) if pool == 'avg': layer.append( nn.AvgPool2d(2, stride=2, ceil_mode=True, count_include_pad=False)) layer.append( init_default( nn.Conv2d(cout, cout * 2, ksize, stride=stride, padding=padding, bias=bias), nn.init.kaiming_normal_)) if use_bn: layer.append(batchnorm_2d(cout * 2, norm_type=bn)) if use_relu: layer.append(relu(True)) if pool == 'max': layer.append(nn.MaxPool2d(2, stride=2)) if pool == 'avg': layer.append( nn.AvgPool2d(2, stride=2, ceil_mode=True, count_include_pad=False)) return nn.Sequential(*layer)