|
|
import torch |
|
|
from torch import nn |
|
|
|
|
|
from stldm.submodules import ChannelConversion |
|
|
from stldm.simvpv2 import stride_generator, ConvSC, MidMetaNet |
|
|
|
|
|
class Encoder(nn.Module): |
|
|
def __init__(self, C_in, C_hid, N_S): |
|
|
super(Encoder, self).__init__() |
|
|
strides = stride_generator(N_S) |
|
|
self.enc = nn.Sequential( |
|
|
ConvSC(C_in, C_hid, stride=strides[0]), |
|
|
*[ConvSC(C_hid, C_hid, stride=s) for s in strides[1:]], |
|
|
ChannelConversion(C_hid, 2*C_hid) |
|
|
) |
|
|
|
|
|
def forward(self, x): |
|
|
for encoder in self.enc: |
|
|
x = encoder(x) |
|
|
(mean, log_var) = torch.chunk(x, 2, dim=1) |
|
|
return mean, log_var |
|
|
|
|
|
class Decoder(nn.Module): |
|
|
def __init__(self, C_hid, C_out, N_S, last_activation='sigmoid'): |
|
|
super(Decoder,self).__init__() |
|
|
strides = stride_generator(N_S, reverse=True) |
|
|
self.dec = nn.Sequential( |
|
|
ChannelConversion(C_hid, C_hid), |
|
|
*[ConvSC(C_hid, C_hid, stride=s, transpose=True) for s in strides[:-1]], |
|
|
ConvSC(C_hid, C_hid, stride=strides[-1], transpose=True) |
|
|
) |
|
|
self.readout = nn.Conv2d(C_hid, C_out, 1) |
|
|
if last_activation=='sigmoid': |
|
|
self.last = nn.Sigmoid() |
|
|
else: |
|
|
self.last = nn.Identity() |
|
|
|
|
|
def forward(self, x): |
|
|
for decoder in self.dec: |
|
|
x = decoder(x) |
|
|
Y = self.readout(x) |
|
|
return self.last(Y) |
|
|
|
|
|
|
|
|
class VAE(nn.Module): |
|
|
def __init__(self, C_in, hid_S, N_S, last_activation='none'): |
|
|
super(VAE, self).__init__() |
|
|
self.encoder = Encoder(C_in, hid_S, N_S) |
|
|
self.decoder = Decoder(hid_S, C_in, N_S, last_activation) |
|
|
|
|
|
def sample_from_standard_normal(self, mean, log_var): |
|
|
std = (0.5 * log_var).exp() |
|
|
return mean + std * torch.randn_like(mean) |
|
|
|
|
|
def encode(self, x): |
|
|
assert x.ndim==4 |
|
|
mean, log_var = self.encoder(x) |
|
|
return mean, log_var |
|
|
|
|
|
def decode(self, z): |
|
|
assert z.ndim==4 |
|
|
dec = self.decoder(z) |
|
|
return dec |
|
|
|
|
|
def kl_from_standard_normal(self, mean, log_var): |
|
|
kl = 0.5 * (log_var.exp() + mean.square() - 1.0 - log_var) |
|
|
return kl.mean() |
|
|
|
|
|
def _losses_(self, x, y): |
|
|
mean, log_var = self.encode(x) |
|
|
kl_loss = self.kl_from_standard_normal(mean, log_var) |
|
|
|
|
|
y_pred = self.forward(x) |
|
|
recon_loss = nn.MSELoss()(y_pred, y) |
|
|
return recon_loss, kl_loss |
|
|
|
|
|
def forward(self, x): |
|
|
mu_z, log_var = self.encode(x) |
|
|
|
|
|
z = self.sample_from_standard_normal(mu_z, log_var) |
|
|
recon = self.decode(z) |
|
|
return recon |
|
|
|
|
|
class SimVPV2_Model(nn.Module): |
|
|
def __init__(self, shape_in, shape_out, hid_S=16, hid_T=256, N_S=4, N_T=4, |
|
|
mlp_ratio=8., drop=0.0, drop_path=0.0, spatio_kernel_enc=3, |
|
|
spatio_kernel_dec=3, last_activation='none', act_inplace=True, **kwargs): |
|
|
super(SimVPV2_Model, self).__init__() |
|
|
T, C, H, W = shape_in |
|
|
T2, C2, H2, W2 = shape_out |
|
|
assert C==C2 and H==H2 and W==W2, 'Need to be the same image shape for input and output' |
|
|
self.T2 = T2 |
|
|
self.T = T |
|
|
|
|
|
H, W = int(H / 2**(N_S/2)), int(W / 2**(N_S/2)) |
|
|
|
|
|
self.vae = VAE(C_in=C, hid_S=hid_S, N_S=N_S, last_activation=last_activation) |
|
|
self.hid = MidMetaNet(T*hid_S, T2*hid_S*2, hid_T, N_T, |
|
|
input_resolution=(H, W), model_type='gsta', |
|
|
mlp_ratio=mlp_ratio, drop=drop, drop_path=drop_path) |
|
|
|
|
|
def forward(self, x_raw): |
|
|
B, T, C, H, W = x_raw.shape |
|
|
x = x_raw.reshape(B*T, C, H, W) |
|
|
|
|
|
embed, log_var = self.vae.encode(x) |
|
|
embed = self.vae.sample_from_standard_normal(embed, log_var) |
|
|
*_, C_, H_, W_ = embed.shape |
|
|
z = embed.view(B, T, C_, H_, W_) |
|
|
|
|
|
hid, *_ = self.hid(z) |
|
|
hid_mu, log_var_hid = torch.chunk(hid, 2, dim=1) |
|
|
hid = self.vae.sample_from_standard_normal(hid_mu, log_var_hid) |
|
|
|
|
|
hid = hid.reshape(B*self.T2, C_, H_, W_) |
|
|
|
|
|
conds_ = hid_mu.reshape(B*self.T2, C_, H_, W_) |
|
|
|
|
|
Y = self.vae.decode(hid) |
|
|
Y = Y.reshape(B, self.T2, C, H, W) |
|
|
return Y, conds_ |
|
|
|
|
|
def _losses_(self, x, y): |
|
|
y_pred, *_ = self.forward(x) |
|
|
recon_loss = nn.MSELoss()(y_pred, y) |
|
|
return recon_loss |
