def __init__(self, cfg, **kwargs): super(MixedActCell, self).__init__() in_channels = get_attr_kwargs(cfg, "in_channels", **kwargs) out_channels = get_attr_kwargs(cfg, "out_channels", **kwargs) ksize = get_attr_kwargs(cfg, "ksize", default=3, **kwargs) up_mode = get_attr_kwargs(cfg, "up_mode", **kwargs) num_skip_in = get_attr_kwargs(cfg, "num_skip_in", default=0, **kwargs) short_cut = get_attr_kwargs(cfg, "short_cut", default=False, **kwargs) norm = get_attr_kwargs(cfg, "norm", default=None, **kwargs) cfg_mix_layer = get_attr_kwargs(cfg, "cfg_mix_layer", **kwargs) cfg_ops = get_attr_kwargs(cfg, "cfg_ops", **kwargs) self.c1 = nn.Conv2d(in_channels, out_channels, ksize, padding=ksize // 2) self.c2 = nn.Conv2d(out_channels, out_channels, ksize, padding=ksize // 2) self.act1 = build_d2layer(cfg_mix_layer, cfg_ops=cfg_ops, num_parameters=in_channels, out_channels=in_channels) self.act2 = build_d2layer(cfg_mix_layer, cfg_ops=cfg_ops, num_parameters=out_channels, out_channels=out_channels) assert up_mode in UP_MODES self.up_mode = up_mode self.norm = norm if norm: assert norm in NORMS if norm == 'bn': self.n1 = nn.BatchNorm2d(in_channels) self.n2 = nn.BatchNorm2d(out_channels) elif norm == 'in': self.n1 = nn.InstanceNorm2d(in_channels) self.n2 = nn.InstanceNorm2d(out_channels) else: raise NotImplementedError(norm) # inner shortcut self.c_sc = None if short_cut: self.c_sc = nn.Conv2d(in_channels, out_channels, kernel_size=1) # cross scale skip self.skip_in_ops = None if num_skip_in: self.skip_in_ops = nn.ModuleList([ nn.Conv2d(out_channels, out_channels, kernel_size=1) for _ in range(num_skip_in) ]) pass
def test_SNConv2d(self): """ """ if 'CUDA_VISIBLE_DEVICES' not in os.environ: os.environ['CUDA_VISIBLE_DEVICES'] = '0' if 'PORT' not in os.environ: os.environ['PORT'] = '6006' if 'TIME_STR' not in os.environ: os.environ['TIME_STR'] = '0' if utils.is_debugging() else '1' # func name assert sys._getframe().f_code.co_name.startswith('test_') command = sys._getframe().f_code.co_name[5:] class_name = self.__class__.__name__[7:] \ if self.__class__.__name__.startswith('Testing') \ else self.__class__.__name__ outdir = f'results/{class_name}/{command}' import yaml from template_lib.d2.layers import build_d2layer cfg_str = """ SNConv2d_3x3: name: "SNConv2d" in_channels: "kwargs['in_channels']" out_channels: "kwargs['out_channels']" kernel_size: 3 padding: 1 """ cfg = EasyDict(yaml.safe_load(cfg_str)) op = build_d2layer(cfg.SNConv2d_3x3, in_channels=8, out_channels=8) op.cuda() x = torch.randn(2, 8, 32, 32).cuda() y = op(x) pass
def test_StyleV2Conv(self): """ """ if 'CUDA_VISIBLE_DEVICES' not in os.environ: os.environ['CUDA_VISIBLE_DEVICES'] = '0' if 'PORT' not in os.environ: os.environ['PORT'] = '6006' if 'TIME_STR' not in os.environ: os.environ['TIME_STR'] = '0' if utils.is_debugging() else '1' # func name assert sys._getframe().f_code.co_name.startswith('test_') command = sys._getframe().f_code.co_name[5:] class_name = self.__class__.__name__[7:] \ if self.__class__.__name__.startswith('Testing') \ else self.__class__.__name__ outdir = f'results/{class_name}/{command}' import yaml from template_lib.d2.layers import build_d2layer cfg_str = """ name: "StyleV2Conv" update_cfg: true """ cfg = EasyDict(yaml.safe_load(cfg_str)) op = build_d2layer(cfg, in_channels=256, out_channels=256) op.cuda() out = op.test_case(in_channels=256, out_channels=256) import torchviz g = torchviz.make_dot(out) g.view() pass
def test_CondInstanceNorm2d(self): """ """ if 'CUDA_VISIBLE_DEVICES' not in os.environ: os.environ['CUDA_VISIBLE_DEVICES'] = '0' if 'PORT' not in os.environ: os.environ['PORT'] = '6006' if 'TIME_STR' not in os.environ: os.environ['TIME_STR'] = '0' if utils.is_debugging() else '1' # func name assert sys._getframe().f_code.co_name.startswith('test_') command = sys._getframe().f_code.co_name[5:] class_name = self.__class__.__name__[7:] \ if self.__class__.__name__.startswith('Testing') \ else self.__class__.__name__ outdir = f'results/{class_name}/{command}' import yaml from template_lib.d2.layers import build_d2layer cfg_str = """ name: "CondInstanceNorm2d" in_features: "kwargs['in_features']" out_features: "kwargs['out_features']" cfg_fc: name: "Linear" in_features: "kwargs['in_features']" out_features: "kwargs['out_features']" """ cfg = EasyDict(yaml.safe_load(cfg_str)) bs = 2 in_features = out_features = 8 op = build_d2layer(cfg, in_features=in_features, out_features=out_features) op.cuda() x = torch.randn(bs, in_features, 32, 32).cuda() y = torch.randn(bs, in_features).cuda() x = op(x, y) import torchviz g = torchviz.make_dot(x) g.view() pass
def test_DepthwiseSeparableConv2d(self): """ """ if 'CUDA_VISIBLE_DEVICES' not in os.environ: os.environ['CUDA_VISIBLE_DEVICES'] = '0' if 'PORT' not in os.environ: os.environ['PORT'] = '6006' if 'TIME_STR' not in os.environ: os.environ['TIME_STR'] = '0' if utils.is_debugging() else '1' # func name assert sys._getframe().f_code.co_name.startswith('test_') command = sys._getframe().f_code.co_name[5:] class_name = self.__class__.__name__[7:] \ if self.__class__.__name__.startswith('Testing') \ else self.__class__.__name__ outdir = f'results/{class_name}/{command}' import yaml from template_lib.d2.layers import build_d2layer cfg_str = """ name: "DepthwiseSeparableConv2d" in_channels: 256 out_channels: 256 kernel_size: 7 padding: 3 """ cfg = EasyDict(yaml.safe_load(cfg_str)) op = build_d2layer(cfg) op.cuda() x = torch.randn(2, 256, 32, 32).cuda() y = op(x) import torchviz g = torchviz.make_dot(y) g.view() pass
def test_MixedLayerCond(self): """ """ if 'CUDA_VISIBLE_DEVICES' not in os.environ: os.environ['CUDA_VISIBLE_DEVICES'] = '0' if 'PORT' not in os.environ: os.environ['PORT'] = '6006' if 'TIME_STR' not in os.environ: os.environ['TIME_STR'] = '0' if utils.is_debugging() else '1' # func name assert sys._getframe().f_code.co_name.startswith('test_') command = sys._getframe().f_code.co_name[5:] class_name = self.__class__.__name__[7:] \ if self.__class__.__name__.startswith('Testing') \ else self.__class__.__name__ outdir = f'results/{class_name}/{command}' import yaml from template_lib.d2.layers import build_d2layer cfg_str = """ layer: name: "MixedLayerCond" in_channels: "kwargs['in_channels']" out_channels: "kwargs['out_channels']" cfg_ops: "kwargs['cfg_ops']" cfg_bn: name: "BatchNorm2d" num_features: "kwargs['num_features']" affine: true track_running_stats: true cfg_act: name: "ReLU" cfg_ops: SNConv2d_3x3: name: "SNConv2d" in_channels: "kwargs['in_channels']" out_channels: "kwargs['out_channels']" kernel_size: 3 padding: 1 Conv2d_3x3: name: "Conv2d" in_channels: "kwargs['in_channels']" out_channels: "kwargs['out_channels']" kernel_size: 3 padding: 1 """ cfg = EasyDict(yaml.safe_load(cfg_str)) op = build_d2layer(cfg.layer, in_channels=8, out_channels=8, cfg_ops=cfg.cfg_ops) num_classes = 2 bs = num_classes num_ops = 2 op.cuda() x = torch.randn(bs, 8, 32, 32).cuda() y = torch.arange(bs).cuda() sample_arc = torch.arange(num_ops).cuda() x = op(x, y, sample_arc) pass
def __init__(self, cfg, **kwargs): super().__init__() self.n_classes = get_attr_kwargs(cfg, 'n_classes', **kwargs) self.ch = get_attr_kwargs(cfg, 'ch', default=512, **kwargs) self.linear_ch = get_attr_kwargs(cfg, 'linear_ch', default=128, **kwargs) self.bottom_width = get_attr_kwargs(cfg, 'bottom_width', default=4, **kwargs) self.dim_z = get_attr_kwargs(cfg, 'dim_z', default=128, **kwargs) self.init_type = get_attr_kwargs(cfg, 'init_type', default='xavier_uniform', **kwargs) self.embedding_dim = get_attr_kwargs(cfg, 'embedding_dim', default=128, **kwargs) self.cfg_upsample = get_attr_kwargs(cfg, 'cfg_upsample', **kwargs) self.num_cells = get_attr_kwargs(cfg, 'num_cells', **kwargs) self.cfg_cell = get_attr_kwargs(cfg, 'cfg_cell', **kwargs) self.cfg_ops = get_attr_kwargs(cfg, 'cfg_ops', **kwargs) self.cfg_out_bn = get_attr_kwargs(cfg, 'cfg_out_bn', **kwargs) self.fixed_arc_file = get_attr_kwargs(cfg, 'fixed_arc_file', default=None, **kwargs) self.fixed_epoch = get_attr_kwargs(cfg, 'fixed_epoch', default=0, **kwargs) self.layer_op_idx = get_attr_kwargs(cfg, 'layer_op_idx', default=None, **kwargs) self.device = torch.device(f'cuda:{comm.get_rank()}') if self.layer_op_idx is not None: # "[1 1 0 1 2 0 2 1 2 0 2 2 2 2 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2]" self.layer_op_idx = json.loads( self.layer_op_idx.replace(' ', ', ').replace('][', '], [').strip()) assert len(self.layer_op_idx) % self.num_cells == 0 num_edges_of_cell = len(self.layer_op_idx) // self.num_cells elif self.fixed_arc_file is not None: sample_arc = self._get_arc_from_file( fixed_arc_file=self.fixed_arc_file, fixed_epoch=self.fixed_epoch, nrows=1) self.layer_op_idx = sample_arc.reshape(-1) assert len(self.layer_op_idx) % self.num_cells == 0 num_edges_of_cell = len(self.layer_op_idx) // self.num_cells self.num_slots = self.num_cells + 1 self.z_chunk_size = (self.dim_z // self.num_slots) self.dim_z_input = self.dim_z self.cbn_in_features = self.embedding_dim + self.z_chunk_size # Prepare class embedding self.class_embedding = nn.Embedding(self.n_classes, self.embedding_dim) self.l1 = nn.Linear(self.dim_z_input, (self.bottom_width**2) * self.linear_ch) self.conv1 = nn.Conv2d(in_channels=self.linear_ch, out_channels=self.ch, kernel_size=1, stride=1, padding=0) self.upsample = build_d2layer(self.cfg_upsample) self.cells = nn.ModuleList() for i in range(self.num_cells): if self.layer_op_idx is not None: cell = build_d2layer( self.cfg_cell, in_channels=self.ch, cfg_ops=self.cfg_ops, cell_op_idx=self.layer_op_idx[i * num_edges_of_cell:(i + 1) * num_edges_of_cell]) else: cell = build_d2layer(self.cfg_cell, in_channels=self.ch, cfg_ops=self.cfg_ops) self.cells.append(cell) self.num_branches = len(self.cfg_ops) self.num_edges_of_cell = cell.num_edges self.num_layers = len(self.cells) * self.num_edges_of_cell out_bn = build_d2layer(self.cfg_out_bn, num_features=self.ch, **kwargs) self.to_rgb = nn.Sequential(out_bn, nn.ReLU(), nn.Conv2d(self.ch, 3, 3, 1, 1), nn.Tanh()) weights_init_func = functools.partial(self.weights_init, init_type=self.init_type) self.apply(weights_init_func)
def __init__(self, cfg, **kwargs): super(DenseGenerator_v2, self).__init__() self.ch = get_attr_kwargs(cfg, 'ch', default=256, **kwargs) self.linear_ch = get_attr_kwargs(cfg, 'linear_ch', default=128, **kwargs) self.bottom_width = get_attr_kwargs(cfg, 'bottom_width', default=4, **kwargs) self.dim_z = get_attr_kwargs(cfg, 'dim_z', default=128, **kwargs) self.init_type = get_attr_kwargs(cfg, 'init_type', default='xavier_uniform', **kwargs) self.cfg_upsample = get_attr_kwargs(cfg, 'cfg_upsample', **kwargs) self.num_cells = get_attr_kwargs(cfg, 'num_cells', **kwargs) self.cfg_cell = get_attr_kwargs(cfg, 'cfg_cell', **kwargs) self.layer_op_idx = get_attr_kwargs(cfg, 'layer_op_idx', default=None, **kwargs) self.cfg_ops = get_attr_kwargs(cfg, 'cfg_ops', **kwargs) self.cfg_out_bn = get_attr_kwargs(cfg, 'cfg_out_bn', **kwargs) if self.layer_op_idx is not None: self.layer_op_idx = json.loads( self.layer_op_idx.replace(' ', ', ').replace('][', '], [').strip()) assert len(self.layer_op_idx) % self.num_cells == 0 num_edges_of_cell = len(self.layer_op_idx) // self.num_cells self.device = torch.device(f'cuda:{comm.get_rank()}') self.l1 = nn.Linear(self.dim_z, (self.bottom_width**2) * self.linear_ch) # self.act = nn.ReLU() self.conv1 = nn.Conv2d(in_channels=self.linear_ch, out_channels=self.ch, kernel_size=1, stride=1, padding=0) self.upsample = build_d2layer(self.cfg_upsample) self.cells = nn.ModuleList() for i in range(self.num_cells): if self.layer_op_idx is not None: cell = build_d2layer( self.cfg_cell, in_channels=self.ch, cfg_ops=self.cfg_ops, cell_op_idx=self.layer_op_idx[i * num_edges_of_cell:(i + 1) * num_edges_of_cell]) else: cell = build_d2layer(self.cfg_cell, in_channels=self.ch, cfg_ops=self.cfg_ops) self.cells.append(cell) self.num_branches = len(self.cfg_ops) self.num_edges_of_cell = cell.num_edges self.num_layers = len(self.cells) * self.num_edges_of_cell out_bn = build_d2layer(self.cfg_out_bn, num_features=self.ch, **kwargs) self.to_rgb = nn.Sequential(out_bn, nn.ReLU(), nn.Conv2d(self.ch, 3, 3, 1, 1), nn.Tanh()) weights_init_func = functools.partial(self.weights_init, init_type=self.init_type) self.apply(weights_init_func)
def __init__(self, cfg, **kwargs): super().__init__() cfg = self.update_cfg(cfg) self.ch = get_attr_kwargs(cfg, 'ch', default=512, **kwargs) self.d_spectral_norm = get_attr_kwargs(cfg, 'd_spectral_norm', default=True, **kwargs) self.init_type = get_attr_kwargs(cfg, 'init_type', default='xavier_uniform', **kwargs) self.cfg_downsample = get_attr_kwargs(cfg, 'cfg_downsample', **kwargs) self.num_cells = get_attr_kwargs(cfg, 'num_cells', **kwargs) self.cfg_cell = get_attr_kwargs(cfg, 'cfg_cell', **kwargs) self.cfg_ops = get_attr_kwargs(cfg, 'cfg_ops', **kwargs) self.fixed_arc_file = get_attr_kwargs(cfg, 'fixed_arc_file', default=None, **kwargs) self.fixed_epoch = get_attr_kwargs(cfg, 'fixed_epoch', default=0, **kwargs) self.layer_op_idx = get_attr_kwargs(cfg, 'layer_op_idx', default=None, **kwargs) self.device = torch.device(f'cuda:{comm.get_rank()}') if self.layer_op_idx is not None: # "[1 1 0 1 2 0 2 1 2 0 2 2 2 2 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2]" self.layer_op_idx = json.loads( self.layer_op_idx.replace(' ', ', ').replace('][', '], [').strip()) assert len(self.layer_op_idx) % self.num_cells == 0 num_edges_of_cell = len(self.layer_op_idx) // self.num_cells elif self.fixed_arc_file is not None: sample_arc = self._get_arc_from_file( fixed_arc_file=self.fixed_arc_file, fixed_epoch=self.fixed_epoch, nrows=1) self.layer_op_idx = sample_arc.reshape(-1) assert len(self.layer_op_idx) % self.num_cells == 0 num_edges_of_cell = len(self.layer_op_idx) // self.num_cells self.conv1 = nn.Conv2d(in_channels=3, out_channels=self.ch, kernel_size=1) if self.d_spectral_norm: self.conv1 = nn.utils.spectral_norm(self.conv1) self.downsample = build_d2layer(self.cfg_downsample) self.cells = nn.ModuleList() for i in range(self.num_cells): if self.layer_op_idx is not None: cell = build_d2layer( self.cfg_cell, in_channels=self.ch, cfg_ops=self.cfg_ops, cell_op_idx=self.layer_op_idx[i * num_edges_of_cell:(i + 1) * num_edges_of_cell]) else: cell = build_d2layer(self.cfg_cell, in_channels=self.ch, cfg_ops=self.cfg_ops) self.cells.append(cell) self.num_branches = len(self.cfg_ops) self.num_edges_of_cell = cell.num_edges self.num_layers = len(self.cells) * self.num_edges_of_cell self.fc = nn.Linear(self.ch, 1, bias=False) if self.d_spectral_norm: self.fc = nn.utils.spectral_norm(self.fc) weights_init_func = functools.partial(self.weights_init, init_type=self.init_type) self.apply(weights_init_func)