First commit OCR_earsing and Synthetics Handwritten Recognition awesome repo
This commit is contained in:
404
OCR_earsing/latent_diffusion/taming/models/vqgan.py
Normal file
404
OCR_earsing/latent_diffusion/taming/models/vqgan.py
Normal file
@@ -0,0 +1,404 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import pytorch_lightning as pl
|
||||
|
||||
from main import instantiate_from_config
|
||||
|
||||
from taming.modules.diffusionmodules.model import Encoder, Decoder
|
||||
from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
|
||||
from taming.modules.vqvae.quantize import GumbelQuantize
|
||||
from taming.modules.vqvae.quantize import EMAVectorQuantizer
|
||||
|
||||
class VQModel(pl.LightningModule):
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
lossconfig,
|
||||
n_embed,
|
||||
embed_dim,
|
||||
ckpt_path=None,
|
||||
ignore_keys=[],
|
||||
image_key="image",
|
||||
colorize_nlabels=None,
|
||||
monitor=None,
|
||||
remap=None,
|
||||
sane_index_shape=False, # tell vector quantizer to return indices as bhw
|
||||
):
|
||||
super().__init__()
|
||||
self.image_key = image_key
|
||||
self.encoder = Encoder(**ddconfig)
|
||||
self.decoder = Decoder(**ddconfig)
|
||||
self.loss = instantiate_from_config(lossconfig)
|
||||
self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
|
||||
remap=remap, sane_index_shape=sane_index_shape)
|
||||
self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
|
||||
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
|
||||
if ckpt_path is not None:
|
||||
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
||||
self.image_key = image_key
|
||||
if colorize_nlabels is not None:
|
||||
assert type(colorize_nlabels)==int
|
||||
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
|
||||
if monitor is not None:
|
||||
self.monitor = monitor
|
||||
|
||||
def init_from_ckpt(self, path, ignore_keys=list()):
|
||||
sd = torch.load(path, map_location="cpu")["state_dict"]
|
||||
keys = list(sd.keys())
|
||||
for k in keys:
|
||||
for ik in ignore_keys:
|
||||
if k.startswith(ik):
|
||||
print("Deleting key {} from state_dict.".format(k))
|
||||
del sd[k]
|
||||
self.load_state_dict(sd, strict=False)
|
||||
print(f"Restored from {path}")
|
||||
|
||||
def encode(self, x):
|
||||
h = self.encoder(x)
|
||||
h = self.quant_conv(h)
|
||||
quant, emb_loss, info = self.quantize(h)
|
||||
return quant, emb_loss, info
|
||||
|
||||
def decode(self, quant):
|
||||
quant = self.post_quant_conv(quant)
|
||||
dec = self.decoder(quant)
|
||||
return dec
|
||||
|
||||
def decode_code(self, code_b):
|
||||
quant_b = self.quantize.embed_code(code_b)
|
||||
dec = self.decode(quant_b)
|
||||
return dec
|
||||
|
||||
def forward(self, input):
|
||||
quant, diff, _ = self.encode(input)
|
||||
dec = self.decode(quant)
|
||||
return dec, diff
|
||||
|
||||
def get_input(self, batch, k):
|
||||
x = batch[k]
|
||||
if len(x.shape) == 3:
|
||||
x = x[..., None]
|
||||
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
|
||||
return x.float()
|
||||
|
||||
def training_step(self, batch, batch_idx, optimizer_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
|
||||
if optimizer_idx == 0:
|
||||
# autoencode
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="train")
|
||||
|
||||
self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
return aeloss
|
||||
|
||||
if optimizer_idx == 1:
|
||||
# discriminator
|
||||
discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="train")
|
||||
self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
return discloss
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="val")
|
||||
|
||||
discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="val")
|
||||
rec_loss = log_dict_ae["val/rec_loss"]
|
||||
self.log("val/rec_loss", rec_loss,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
|
||||
self.log("val/aeloss", aeloss,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
|
||||
self.log_dict(log_dict_ae)
|
||||
self.log_dict(log_dict_disc)
|
||||
return self.log_dict
|
||||
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
|
||||
list(self.decoder.parameters())+
|
||||
list(self.quantize.parameters())+
|
||||
list(self.quant_conv.parameters())+
|
||||
list(self.post_quant_conv.parameters()),
|
||||
lr=lr, betas=(0.5, 0.9))
|
||||
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
|
||||
lr=lr, betas=(0.5, 0.9))
|
||||
return [opt_ae, opt_disc], []
|
||||
|
||||
def get_last_layer(self):
|
||||
return self.decoder.conv_out.weight
|
||||
|
||||
def log_images(self, batch, **kwargs):
|
||||
log = dict()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
x = x.to(self.device)
|
||||
xrec, _ = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec.shape[1] > 3
|
||||
x = self.to_rgb(x)
|
||||
xrec = self.to_rgb(xrec)
|
||||
log["inputs"] = x
|
||||
log["reconstructions"] = xrec
|
||||
return log
|
||||
|
||||
def to_rgb(self, x):
|
||||
assert self.image_key == "segmentation"
|
||||
if not hasattr(self, "colorize"):
|
||||
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
|
||||
x = F.conv2d(x, weight=self.colorize)
|
||||
x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
|
||||
return x
|
||||
|
||||
|
||||
class VQSegmentationModel(VQModel):
|
||||
def __init__(self, n_labels, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.register_buffer("colorize", torch.randn(3, n_labels, 1, 1))
|
||||
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
|
||||
list(self.decoder.parameters())+
|
||||
list(self.quantize.parameters())+
|
||||
list(self.quant_conv.parameters())+
|
||||
list(self.post_quant_conv.parameters()),
|
||||
lr=lr, betas=(0.5, 0.9))
|
||||
return opt_ae
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="train")
|
||||
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
return aeloss
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="val")
|
||||
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
total_loss = log_dict_ae["val/total_loss"]
|
||||
self.log("val/total_loss", total_loss,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
|
||||
return aeloss
|
||||
|
||||
@torch.no_grad()
|
||||
def log_images(self, batch, **kwargs):
|
||||
log = dict()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
x = x.to(self.device)
|
||||
xrec, _ = self(x)
|
||||
if x.shape[1] > 3:
|
||||
# colorize with random projection
|
||||
assert xrec.shape[1] > 3
|
||||
# convert logits to indices
|
||||
xrec = torch.argmax(xrec, dim=1, keepdim=True)
|
||||
xrec = F.one_hot(xrec, num_classes=x.shape[1])
|
||||
xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float()
|
||||
x = self.to_rgb(x)
|
||||
xrec = self.to_rgb(xrec)
|
||||
log["inputs"] = x
|
||||
log["reconstructions"] = xrec
|
||||
return log
|
||||
|
||||
|
||||
class VQNoDiscModel(VQModel):
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
lossconfig,
|
||||
n_embed,
|
||||
embed_dim,
|
||||
ckpt_path=None,
|
||||
ignore_keys=[],
|
||||
image_key="image",
|
||||
colorize_nlabels=None
|
||||
):
|
||||
super().__init__(ddconfig=ddconfig, lossconfig=lossconfig, n_embed=n_embed, embed_dim=embed_dim,
|
||||
ckpt_path=ckpt_path, ignore_keys=ignore_keys, image_key=image_key,
|
||||
colorize_nlabels=colorize_nlabels)
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
# autoencode
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="train")
|
||||
output = pl.TrainResult(minimize=aeloss)
|
||||
output.log("train/aeloss", aeloss,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
output.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
return output
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="val")
|
||||
rec_loss = log_dict_ae["val/rec_loss"]
|
||||
output = pl.EvalResult(checkpoint_on=rec_loss)
|
||||
output.log("val/rec_loss", rec_loss,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
output.log("val/aeloss", aeloss,
|
||||
prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
output.log_dict(log_dict_ae)
|
||||
|
||||
return output
|
||||
|
||||
def configure_optimizers(self):
|
||||
optimizer = torch.optim.Adam(list(self.encoder.parameters())+
|
||||
list(self.decoder.parameters())+
|
||||
list(self.quantize.parameters())+
|
||||
list(self.quant_conv.parameters())+
|
||||
list(self.post_quant_conv.parameters()),
|
||||
lr=self.learning_rate, betas=(0.5, 0.9))
|
||||
return optimizer
|
||||
|
||||
|
||||
class GumbelVQ(VQModel):
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
lossconfig,
|
||||
n_embed,
|
||||
embed_dim,
|
||||
temperature_scheduler_config,
|
||||
ckpt_path=None,
|
||||
ignore_keys=[],
|
||||
image_key="image",
|
||||
colorize_nlabels=None,
|
||||
monitor=None,
|
||||
kl_weight=1e-8,
|
||||
remap=None,
|
||||
):
|
||||
|
||||
z_channels = ddconfig["z_channels"]
|
||||
super().__init__(ddconfig,
|
||||
lossconfig,
|
||||
n_embed,
|
||||
embed_dim,
|
||||
ckpt_path=None,
|
||||
ignore_keys=ignore_keys,
|
||||
image_key=image_key,
|
||||
colorize_nlabels=colorize_nlabels,
|
||||
monitor=monitor,
|
||||
)
|
||||
|
||||
self.loss.n_classes = n_embed
|
||||
self.vocab_size = n_embed
|
||||
|
||||
self.quantize = GumbelQuantize(z_channels, embed_dim,
|
||||
n_embed=n_embed,
|
||||
kl_weight=kl_weight, temp_init=1.0,
|
||||
remap=remap)
|
||||
|
||||
self.temperature_scheduler = instantiate_from_config(temperature_scheduler_config) # annealing of temp
|
||||
|
||||
if ckpt_path is not None:
|
||||
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
||||
|
||||
def temperature_scheduling(self):
|
||||
self.quantize.temperature = self.temperature_scheduler(self.global_step)
|
||||
|
||||
def encode_to_prequant(self, x):
|
||||
h = self.encoder(x)
|
||||
h = self.quant_conv(h)
|
||||
return h
|
||||
|
||||
def decode_code(self, code_b):
|
||||
raise NotImplementedError
|
||||
|
||||
def training_step(self, batch, batch_idx, optimizer_idx):
|
||||
self.temperature_scheduling()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x)
|
||||
|
||||
if optimizer_idx == 0:
|
||||
# autoencode
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="train")
|
||||
|
||||
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
self.log("temperature", self.quantize.temperature, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
return aeloss
|
||||
|
||||
if optimizer_idx == 1:
|
||||
# discriminator
|
||||
discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="train")
|
||||
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
|
||||
return discloss
|
||||
|
||||
def validation_step(self, batch, batch_idx):
|
||||
x = self.get_input(batch, self.image_key)
|
||||
xrec, qloss = self(x, return_pred_indices=True)
|
||||
aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="val")
|
||||
|
||||
discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step,
|
||||
last_layer=self.get_last_layer(), split="val")
|
||||
rec_loss = log_dict_ae["val/rec_loss"]
|
||||
self.log("val/rec_loss", rec_loss,
|
||||
prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
|
||||
self.log("val/aeloss", aeloss,
|
||||
prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
|
||||
self.log_dict(log_dict_ae)
|
||||
self.log_dict(log_dict_disc)
|
||||
return self.log_dict
|
||||
|
||||
def log_images(self, batch, **kwargs):
|
||||
log = dict()
|
||||
x = self.get_input(batch, self.image_key)
|
||||
x = x.to(self.device)
|
||||
# encode
|
||||
h = self.encoder(x)
|
||||
h = self.quant_conv(h)
|
||||
quant, _, _ = self.quantize(h)
|
||||
# decode
|
||||
x_rec = self.decode(quant)
|
||||
log["inputs"] = x
|
||||
log["reconstructions"] = x_rec
|
||||
return log
|
||||
|
||||
|
||||
class EMAVQ(VQModel):
|
||||
def __init__(self,
|
||||
ddconfig,
|
||||
lossconfig,
|
||||
n_embed,
|
||||
embed_dim,
|
||||
ckpt_path=None,
|
||||
ignore_keys=[],
|
||||
image_key="image",
|
||||
colorize_nlabels=None,
|
||||
monitor=None,
|
||||
remap=None,
|
||||
sane_index_shape=False, # tell vector quantizer to return indices as bhw
|
||||
):
|
||||
super().__init__(ddconfig,
|
||||
lossconfig,
|
||||
n_embed,
|
||||
embed_dim,
|
||||
ckpt_path=None,
|
||||
ignore_keys=ignore_keys,
|
||||
image_key=image_key,
|
||||
colorize_nlabels=colorize_nlabels,
|
||||
monitor=monitor,
|
||||
)
|
||||
self.quantize = EMAVectorQuantizer(n_embed=n_embed,
|
||||
embedding_dim=embed_dim,
|
||||
beta=0.25,
|
||||
remap=remap)
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
#Remove self.quantize from parameter list since it is updated via EMA
|
||||
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
|
||||
list(self.decoder.parameters())+
|
||||
list(self.quant_conv.parameters())+
|
||||
list(self.post_quant_conv.parameters()),
|
||||
lr=lr, betas=(0.5, 0.9))
|
||||
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
|
||||
lr=lr, betas=(0.5, 0.9))
|
||||
return [opt_ae, opt_disc], []
|
||||
Reference in New Issue
Block a user