Generator Class Reference

Inheritance diagram for Generator:

Collaboration diagram for Generator:

Public Member Functions
def	__init__ (self, G_ch=64, G_depth=2, dim_z=128, bottom_width=4, resolution=256, G_kernel_size=3, G_attn="64", n_classes=40, H_base=1, num_G_SVs=1, num_G_SV_itrs=1, attn_type="sa", G_shared=True, shared_dim=128, hier=True, cross_replica=False, mybn=False, G_activation=nn.ReLU(inplace=False), G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8, BN_eps=1e-5, SN_eps=1e-12, G_init="ortho", G_mixed_precision=False, G_fp16=False, skip_init=False, no_optim=False, sched_version="default", RRM_prx_G=True, n_head_G=2, G_param="SN", norm_style="bn", **kwargs)
	Constructor.
def	init_weights (self)
	Initialize.
def	forward (self, z, y)
	forward

Public Attributes
	ch
	Channel width mulitplier.
	G_depth
	Number of resblocks per stage.
	dim_z
	Dimensionality of the latent space.
	bottom_width
	The initial spatial dimensions.
	H_base
	The initial harizontal dimension.
	resolution
	Resolution of the output.
	kernel_size
	Kernel size?
	attention
	Attention?
	n_classes
	number of classes, for use in categorical conditional generation
	G_shared
	Use shared embeddings?
	shared_dim
	Dimensionality of the shared embedding? Unused if not using G_shared.
	hier
	Hierarchical latent space?
	cross_replica
	Cross replica batchnorm?
	mybn
	Use my batchnorm?
	activation
	activation
	init
	Initialization style.
	G_param
	Parameterization style.
	norm_style
	Normalization style.
	BN_eps
	Epsilon for BatchNorm?
	SN_eps
	Epsilon for Spectral Norm?
	fp16
	fp16?
	arch
	Architecture dict.
	RRM_prx_G
	RRM_prx_G.
	n_head_G
	n_head_G
	which_conv
	which conv
	which_linear
	which linear
	which_embedding
	which embedding
	which_bn
	which bn
	shared
	shared
	RR_G
	RRM on proxy embeddings.
	linear
	First linear layer.
	blocks
	blocks
	output_layer
	output layer
	lr
	lr
	B1
	B1.
	B2
	B2.
	adam_eps
	adam_eps
	optim
	optim
	lr_sched
	lr sched
	param_count
	parameter count

Detailed Description

Generator

Definition at line 1127 of file ieagan.py.

Constructor & Destructor Documentation

__init__()

def __init__	(		self,
			G_ch = 64,
			G_depth = 2,
			dim_z = 128,
			bottom_width = 4,
			resolution = 256,
			G_kernel_size = 3,
			G_attn = "64",
			n_classes = 40,
			H_base = 1,
			num_G_SVs = 1,
			num_G_SV_itrs = 1,
			attn_type = "sa",
			G_shared = True,
			shared_dim = 128,
			hier = True,
			cross_replica = False,
			mybn = False,
			G_activation = nn.ReLU(inplace=False),
			G_lr = 5e-5,
			G_B1 = 0.0,
			G_B2 = 0.999,
			adam_eps = 1e-8,
			BN_eps = 1e-5,
			SN_eps = 1e-12,
			G_init = "ortho",
			G_mixed_precision = False,
			G_fp16 = False,
			skip_init = False,
			no_optim = False,
			sched_version = "default",
			RRM_prx_G = True,
			n_head_G = 2,
			G_param = "SN",
			norm_style = "bn",
		**	kwargs
	)

Constructor.

Reimplemented in Model.

Definition at line 1131 of file ieagan.py.

    ):
        super(Generator, self).__init__()
        
        self.ch = G_ch
        
        self.G_depth = G_depth
        
        self.dim_z = dim_z
        
        self.bottom_width = bottom_width
        
        self.H_base = H_base
        
        self.resolution = resolution
        
        self.kernel_size = G_kernel_size
        
        self.attention = G_attn
        
        self.n_classes = n_classes
        
        self.G_shared = G_shared
        
        self.shared_dim = shared_dim if shared_dim > 0 else dim_z
        
        self.hier = hier
        
        self.cross_replica = cross_replica
        
        self.mybn = mybn
        # nonlinearity for residual blocks
        if G_activation == "inplace_relu":
            
            self.activation = torch.nn.ReLU(inplace=True)
        elif G_activation == "relu":
            self.activation = torch.nn.ReLU(inplace=False)
        elif G_activation == "leaky_relu":
            self.activation = torch.nn.LeakyReLU(0.2, inplace=False)
        else:
            raise NotImplementedError("activation function not implemented")
        
        self.init = G_init
        
        self.G_param = G_param
        
        self.norm_style = norm_style
        
        self.BN_eps = BN_eps
        
        self.SN_eps = SN_eps
        
        self.fp16 = G_fp16
        
        self.arch = G_arch(self.ch, self.attention)[resolution]
        
        self.RRM_prx_G = RRM_prx_G
        
        self.n_head_G = n_head_G
 
        # Which convs, batchnorms, and linear layers to use
        if self.G_param == "SN":
            
            self.which_conv = functools.partial(
                SNConv2d,
                kernel_size=3,
                padding=1,
                num_svs=num_G_SVs,
                num_itrs=num_G_SV_itrs,
                eps=self.SN_eps,
            )
            
            self.which_linear = functools.partial(
                SNLinear,
                num_svs=num_G_SVs,
                num_itrs=num_G_SV_itrs,
                eps=self.SN_eps,
            )
        else:
            self.which_conv = functools.partial(nn.Conv2d, kernel_size=3, padding=1)
            self.which_linear = nn.Linear
 
        # We use a non-spectral-normed embedding here regardless;
        # For some reason applying SN to G's embedding seems to randomly cripple G  # noqa
        
        self.which_embedding = nn.Embedding
        bn_linear = (
            functools.partial(self.which_linear, bias=False)
            if self.G_shared
            else self.which_embedding
        )
        
        self.which_bn = functools.partial(
            ccbn,
            which_linear=bn_linear,
            cross_replica=self.cross_replica,
            mybn=self.mybn,
            input_size=(
                self.shared_dim + self.dim_z if self.G_shared else self.n_classes
            ),
            norm_style=self.norm_style,
            eps=self.BN_eps,
        )
        
        self.shared = (
            self.which_embedding(n_classes, self.shared_dim)
            if G_shared
            else identity()
        )
 
        if self.RRM_prx_G:
            
            self.RR_G = RelationalReasoning(
                num_layers=1,
                input_dim=128,
                dim_feedforward=128,
                which_linear=nn.Linear,
                num_heads=self.n_head_G,
                dropout=0.0,
                hidden_dim=128,
            )
 
        
        self.linear = self.which_linear(
            self.dim_z + self.shared_dim,
            self.arch["in_channels"][0] * ((self.bottom_width**2) * self.H_base),
        )
 
        # self.blocks is a doubly-nested list of modules, the outer loop intended  # noqa
        # to be over blocks at a given resolution (resblocks and/or self-attention)  # noqa
        # while the inner loop is over a given block
        
        self.blocks = []
        for index in range(len(self.arch["out_channels"])):
            self.blocks += [
                [
                    GBlock(
                        in_channels=self.arch["in_channels"][index],
                        out_channels=self.arch["in_channels"][index]
                        if g_index == 0
                        else self.arch["out_channels"][index],
                        which_conv=self.which_conv,
                        which_bn=self.which_bn,
                        activation=self.activation,
                        upsample=(
                            functools.partial(F.interpolate, scale_factor=2)
                            if self.arch["upsample"][index]
                            and g_index == (self.G_depth - 1)
                            else None
                        ),
                    )
                ]
                for g_index in range(self.G_depth)
            ]
 
            # If attention on this block, attach it to the end
            if self.arch["attention"][self.arch["resolution"][index]]:
                print(
                    f"Adding attention layer in G at resolution {self.arch['resolution'][index]:d}"
                )
 
                if attn_type == "sa":
                    self.blocks[-1] += [
                        Attention(self.arch["out_channels"][index], self.which_conv)
                    ]
                elif attn_type == "cbam":
                    self.blocks[-1] += [
                        CBAM_attention(
                            self.arch["out_channels"][index], self.which_conv
                        )
                    ]
                elif attn_type == "ila":
                    self.blocks[-1] += [ILA(self.arch["out_channels"][index])]
 
        # Turn self.blocks into a ModuleList so that it's all properly registered.  # noqa
        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
 
        # output layer: batchnorm-relu-conv.
        # Consider using a non-spectral conv here
        
        self.output_layer = nn.Sequential(
            bn(
                self.arch["out_channels"][-1],
                cross_replica=self.cross_replica,
                mybn=self.mybn,
            ),
            self.activation,
            self.which_conv(self.arch["out_channels"][-1], 1),
        )
 
        # Initialize weights. Optionally skip init for testing.
        if not skip_init:
            self.init_weights()
 
        # Set up optimizer
        # If this is an EMA copy, no need for an optim, so just return now
        if no_optim:
            return
        
        self.lr = G_lr
        
        self.B1 = G_B1
        
        self.B2 = G_B2
        
        self.adam_eps = adam_eps
        if G_mixed_precision:
            print("Using fp16 adam in G...")
            import utils
 
            self.optim = utils.Adam16(
                params=self.parameters(),
                lr=self.lr,
                betas=(self.B1, self.B2),
                weight_decay=0,
                eps=self.adam_eps,
            )
 
        
        self.optim = optim.Adam(
            params=self.parameters(),
            lr=self.lr,
            betas=(self.B1, self.B2),
            weight_decay=0,
            eps=self.adam_eps,
        )
        # LR scheduling
        if sched_version == "default":
            
            self.lr_sched = None
        elif sched_version == "CosAnnealLR":
            self.lr_sched = optim.lr_scheduler.CosineAnnealingLR(
                self.optim,
                T_max=kwargs["num_epochs"],
                eta_min=self.lr / 4,
                last_epoch=-1,
            )
        elif sched_version == "CosAnnealWarmRes":
            self.lr_sched = optim.lr_scheduler.CosineAnnealingWarmRestarts(
                self.optim, T_0=10, T_mult=2, eta_min=self.lr / 4
            )
        else:
            self.lr_sched = None
 

Member Function Documentation

forward()

def forward	(		self,
			z,
			y
	)

forward

Definition at line 1435 of file ieagan.py.

    def forward(self, z, y):
        y = self.shared(y)
        # If relational embedding
        if self.RRM_prx_G:
            y = self.RR_G(y.unsqueeze(0)).squeeze(0)
            # y = F.normalize(y, dim=1)
        # If hierarchical, concatenate zs and ys
        if self.hier:  # y and z are [bs,128] dimensional
            z = torch.cat([y, z], 1)
            y = z
        # First linear layer
        h = self.linear(z)  # ([bs,256]-->[bs,24576])
        # Reshape
        h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width * self.H_base)
        # Loop over blocks
        for _, blocklist in enumerate(self.blocks):
            # Second inner loop in case block has multiple layers
            for block in blocklist:
                h = block(h, y)
 
        # Apply batchnorm-relu-conv-tanh at output
        return torch.tanh(self.output_layer(h))
 
 

init_weights()

def init_weights

(

self

)

Initialize.

Definition at line 1412 of file ieagan.py.

    def init_weights(self):
        
        self.param_count = 0
        for module in self.modules():
            if (
                isinstance(module, nn.Conv2d)
                or isinstance(module, nn.Linear)
                or isinstance(module, nn.Embedding)
            ):
                if self.init == "ortho":
                    init.orthogonal_(module.weight)
                elif self.init == "N02":
                    init.normal_(module.weight, 0, 0.02)
                elif self.init in ["glorot", "xavier"]:
                    init.xavier_uniform_(module.weight)
                else:
                    print("Init style not recognized...")
                self.param_count += sum(
                    [p.data.nelement() for p in module.parameters()]
                )
        print(f"Param count for G's initialized parameters: {self.param_count}")
 

Member Data Documentation

activation

activation

activation

Definition at line 1201 of file ieagan.py.

adam_eps

adam_eps

adam_eps

Definition at line 1372 of file ieagan.py.

arch

arch

Architecture dict.

Definition at line 1221 of file ieagan.py.

attention

attention

Attention?

Definition at line 1185 of file ieagan.py.

B1

B1

B1.

Definition at line 1368 of file ieagan.py.

B2

B2

B2.

Definition at line 1370 of file ieagan.py.

blocks

blocks

blocks

Definition at line 1299 of file ieagan.py.

BN_eps

BN_eps

Epsilon for BatchNorm?

Definition at line 1215 of file ieagan.py.

bottom_width

bottom_width

The initial spatial dimensions.

Definition at line 1177 of file ieagan.py.

ch

ch

Channel width mulitplier.

Definition at line 1171 of file ieagan.py.

cross_replica

cross_replica

Cross replica batchnorm?

Definition at line 1195 of file ieagan.py.

dim_z

dim_z

Dimensionality of the latent space.

Definition at line 1175 of file ieagan.py.

fp16

fp16

fp16?

Definition at line 1219 of file ieagan.py.

G_depth

G_depth

Number of resblocks per stage.

Definition at line 1173 of file ieagan.py.

G_param

G_param

Parameterization style.

Definition at line 1211 of file ieagan.py.

G_shared

G_shared

Use shared embeddings?

Definition at line 1189 of file ieagan.py.

H_base

H_base

The initial harizontal dimension.

Definition at line 1179 of file ieagan.py.

hier

hier

Hierarchical latent space?

Definition at line 1193 of file ieagan.py.

init

init

Initialization style.

Definition at line 1209 of file ieagan.py.

kernel_size

kernel_size

Kernel size?

Definition at line 1183 of file ieagan.py.

linear

linear

First linear layer.

Definition at line 1290 of file ieagan.py.

lr

lr

lr

Definition at line 1366 of file ieagan.py.

lr_sched

lr_sched

lr sched

Definition at line 1396 of file ieagan.py.

mybn

mybn

Use my batchnorm?

Definition at line 1197 of file ieagan.py.

n_classes

n_classes

number of classes, for use in categorical conditional generation

Definition at line 1187 of file ieagan.py.

n_head_G

n_head_G

n_head_G

Definition at line 1225 of file ieagan.py.

norm_style

norm_style

Normalization style.

Definition at line 1213 of file ieagan.py.

optim

optim

optim

Definition at line 1377 of file ieagan.py.

output_layer

output_layer

output layer

Definition at line 1347 of file ieagan.py.

param_count

param_count

parameter count

Definition at line 1414 of file ieagan.py.

resolution

resolution

Resolution of the output.

Definition at line 1181 of file ieagan.py.

RR_G

RR_G

RRM on proxy embeddings.

Definition at line 1279 of file ieagan.py.

RRM_prx_G

RRM_prx_G

RRM_prx_G.

Definition at line 1223 of file ieagan.py.

shared

shared

shared

Definition at line 1271 of file ieagan.py.

shared_dim

shared_dim

Dimensionality of the shared embedding? Unused if not using G_shared.

Definition at line 1191 of file ieagan.py.

SN_eps

SN_eps

Epsilon for Spectral Norm?

Definition at line 1217 of file ieagan.py.

which_bn

which_bn

which bn

Definition at line 1259 of file ieagan.py.

which_conv

which_conv

which conv

Definition at line 1230 of file ieagan.py.

which_embedding

which_embedding

which embedding

Definition at line 1252 of file ieagan.py.

which_linear

which_linear

which linear

Definition at line 1239 of file ieagan.py.

---

The documentation for this class was generated from the following file:

• pxd/scripts/pxd/background_generator/models/ieagan.py