338 lines
13 KiB
Python
338 lines
13 KiB
Python
import time
|
|
from contextlib import contextmanager
|
|
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from ..modules.diffusionmodules.model import Decoder, Encoder
|
|
from ..modules.distributions.distributions import DiagonalGaussianDistribution
|
|
from ..util import instantiate_from_config
|
|
from .enc_dec_pytorch import Decoder as Mag_Decoder
|
|
from .enc_dec_pytorch import Encoder as Mag_Encoder
|
|
|
|
|
|
class AutoencoderKLMagvit(pl.LightningModule):
|
|
def __init__(self,
|
|
ddconfig,
|
|
lossconfig,
|
|
embed_dim,
|
|
ckpt_path=None,
|
|
ignore_keys=[],
|
|
image_key="image",
|
|
colorize_nlabels=None,
|
|
monitor=None,
|
|
):
|
|
super().__init__()
|
|
self.image_key = image_key
|
|
self.encoder = Mag_Encoder()
|
|
self.decoder = Mag_Decoder()
|
|
self.loss = instantiate_from_config(lossconfig)
|
|
self.quant_conv = torch.nn.Conv3d(16, 16, 1)
|
|
self.post_quant_conv = torch.nn.Conv3d(8, 8, 1)
|
|
self.embed_dim = embed_dim
|
|
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
|
|
if ckpt_path is not None:
|
|
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
|
|
|
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)
|
|
moments = self.quant_conv(h)
|
|
posterior = DiagonalGaussianDistribution(moments)
|
|
return posterior
|
|
|
|
def decode(self, z):
|
|
z = self.post_quant_conv(z)
|
|
dec = self.decoder(z)
|
|
return dec
|
|
|
|
def forward(self, input, sample_posterior=True):
|
|
if input.ndim==4:
|
|
input = input.unsqueeze(2)
|
|
posterior = self.encode(input)
|
|
if sample_posterior:
|
|
z = posterior.sample()
|
|
else:
|
|
z = posterior.mode()
|
|
dec = self.decode(z)
|
|
return dec, posterior
|
|
|
|
def get_input(self, batch, k):
|
|
x = batch[k]
|
|
if x.ndim==5:
|
|
x = x.permute(0, 4, 1, 2, 3).to(memory_format=torch.contiguous_format).float()
|
|
return x
|
|
if len(x.shape) == 3:
|
|
x = x[..., None]
|
|
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float()
|
|
return x
|
|
|
|
def training_step(self, batch, batch_idx, optimizer_idx):
|
|
# tic = time.time()
|
|
inputs = self.get_input(batch, self.image_key)
|
|
# print(f"get_input time {time.time() - tic}")
|
|
# tic = time.time()
|
|
reconstructions, posterior = self(inputs)
|
|
# print(f"model forward time {time.time() - tic}")
|
|
|
|
if optimizer_idx == 0:
|
|
# train encoder+decoder+logvar
|
|
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
|
|
last_layer=self.get_last_layer(), split="train")
|
|
self.log("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=False)
|
|
# print(f"cal loss time {time.time() - tic}")
|
|
return aeloss
|
|
|
|
if optimizer_idx == 1:
|
|
# train the discriminator
|
|
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
|
|
last_layer=self.get_last_layer(), split="train")
|
|
|
|
self.log("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=False)
|
|
# print(f"cal loss time {time.time() - tic}")
|
|
return discloss
|
|
|
|
def validation_step(self, batch, batch_idx):
|
|
with torch.no_grad():
|
|
inputs = self.get_input(batch, self.image_key)
|
|
reconstructions, posterior = self(inputs)
|
|
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
|
|
last_layer=self.get_last_layer(), split="val")
|
|
|
|
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
|
|
last_layer=self.get_last_layer(), split="val")
|
|
|
|
self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
|
|
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.AdamW(list(self.encoder.parameters())+
|
|
list(self.decoder.parameters())+
|
|
list(self.quant_conv.parameters())+
|
|
list(self.post_quant_conv.parameters()),
|
|
lr=lr, betas=(0.9, 0.999), weight_decay=5e-2)
|
|
opt_disc = torch.optim.AdamW(self.loss.discriminator.parameters(),
|
|
lr=lr, betas=(0.9, 0.999), weight_decay=5e-2)
|
|
return [opt_ae, opt_disc], []
|
|
|
|
def get_last_layer(self):
|
|
return self.decoder.conv_out.weight
|
|
|
|
@torch.no_grad()
|
|
def log_images(self, batch, only_inputs=False, **kwargs):
|
|
log = dict()
|
|
x = self.get_input(batch, self.image_key)
|
|
x = x.to(self.device)
|
|
if not only_inputs:
|
|
xrec, posterior = 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["samples"] = self.decode(torch.randn_like(posterior.sample()))
|
|
log["reconstructions"] = xrec
|
|
log["inputs"] = x
|
|
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 AutoencoderKL(pl.LightningModule):
|
|
def __init__(self,
|
|
ddconfig,
|
|
lossconfig,
|
|
embed_dim,
|
|
ckpt_path=None,
|
|
ignore_keys=[],
|
|
image_key="image",
|
|
colorize_nlabels=None,
|
|
monitor=None,
|
|
):
|
|
super().__init__()
|
|
self.image_key = image_key
|
|
self.encoder = Encoder(**ddconfig)
|
|
self.decoder = Decoder(**ddconfig)
|
|
self.loss = instantiate_from_config(lossconfig)
|
|
assert ddconfig["double_z"]
|
|
self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1)
|
|
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
|
|
self.embed_dim = embed_dim
|
|
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
|
|
if ckpt_path is not None:
|
|
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
|
|
|
|
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)
|
|
moments = self.quant_conv(h)
|
|
posterior = DiagonalGaussianDistribution(moments)
|
|
return posterior
|
|
|
|
def decode(self, z):
|
|
z = self.post_quant_conv(z)
|
|
dec = self.decoder(z)
|
|
return dec
|
|
|
|
def forward(self, input, sample_posterior=True):
|
|
posterior = self.encode(input)
|
|
if sample_posterior:
|
|
z = posterior.sample()
|
|
else:
|
|
z = posterior.mode()
|
|
dec = self.decode(z)
|
|
return dec, posterior
|
|
|
|
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).float()
|
|
return x
|
|
|
|
def training_step(self, batch, batch_idx, optimizer_idx):
|
|
# tic = time.time()
|
|
inputs = self.get_input(batch, self.image_key)
|
|
# print(f"get_input time {time.time() - tic}")
|
|
# tic = time.time()
|
|
reconstructions, posterior = self(inputs)
|
|
# print(f"model forward time {time.time() - tic}")
|
|
tic = time.time()
|
|
|
|
if optimizer_idx == 0:
|
|
# train encoder+decoder+logvar
|
|
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
|
|
last_layer=self.get_last_layer(), split="train")
|
|
self.log("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=False)
|
|
# print(f"cal loss time {time.time() - tic}")
|
|
return aeloss
|
|
|
|
if optimizer_idx == 1:
|
|
# train the discriminator
|
|
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
|
|
last_layer=self.get_last_layer(), split="train")
|
|
|
|
self.log("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=False)
|
|
# print(f"cal loss time {time.time() - tic}")
|
|
return discloss
|
|
|
|
def validation_step(self, batch, batch_idx):
|
|
tic = time.time()
|
|
inputs = self.get_input(batch, self.image_key)
|
|
print(f"get_input time {time.time() - tic}")
|
|
tic = time.time()
|
|
reconstructions, posterior = self(inputs)
|
|
print(f"val forward time {time.time() - tic}")
|
|
tic = time.time()
|
|
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
|
|
last_layer=self.get_last_layer(), split="val")
|
|
|
|
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
|
|
last_layer=self.get_last_layer(), split="val")
|
|
|
|
self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
|
|
self.log_dict(log_dict_ae)
|
|
self.log_dict(log_dict_disc)
|
|
print(f"val end time {time.time() - tic}")
|
|
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.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
|
|
|
|
@torch.no_grad()
|
|
def log_images(self, batch, only_inputs=False, **kwargs):
|
|
log = dict()
|
|
x = self.get_input(batch, self.image_key)
|
|
x = x.to(self.device)
|
|
if not only_inputs:
|
|
xrec, posterior = 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["samples"] = self.decode(torch.randn_like(posterior.sample()))
|
|
log["reconstructions"] = xrec
|
|
log["inputs"] = x
|
|
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 IdentityFirstStage(torch.nn.Module):
|
|
def __init__(self, *args, vq_interface=False, **kwargs):
|
|
self.vq_interface = vq_interface # TODO: Should be true by default but check to not break older stuff
|
|
super().__init__()
|
|
|
|
def encode(self, x, *args, **kwargs):
|
|
return x
|
|
|
|
def decode(self, x, *args, **kwargs):
|
|
return x
|
|
|
|
def quantize(self, x, *args, **kwargs):
|
|
if self.vq_interface:
|
|
return x, None, [None, None, None]
|
|
return x
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
return x
|