AEVB / SGVB를 적용한 Variational Autoencoder를 구현한다.
Model
본래 VAE는 MLP로 구성되지만, 이 구현에선 컨볼루션 레이어를 사용한다. latent_dim은 잠재변수 z의 차원이다. 2와 4에서 실험할 것이다.
Encoder
1 -> 32, size=3x3, stride=2, pad=1
32 -> 64, size=3x3, stride=2, pad=1
정규분포의 분산은 항상 양수인데, linear 레이어의 출력은 실수값이다. 실수 범위를 가지는 로그분산( log σ 2 \log \sigma^2 log σ 2 σ = exp ( 1 2 log σ 2 ) \sigma=\exp(\frac 1 2 \log\sigma^2) σ = exp ( 2 1 log σ 2 )
Sampling
torch.randn_like 로 epsilon 샘플 후, 이를 이용해 z를 계산한다.
Decoder
베르누이 분포를 가정한다. 출력이 0과 1 사이의 픽셀값이다.
64 -> 64, size=3x3, stride=2
64 -> 32, size=3x3, stride=2
32 -> 1, size=4x4, stride=1
ConvTranspose2d로 차원(28x28)을 맟출 수 없었기에 최종적으로 Conv2d를 한번 더 해준다.
모델 클래스의 코드는 다음과 같다. 21/07/26 수정, 모델을 더 유연하게 변환했다.
Copy class ConvVAE(nn.Module):
def __init__(self,input_size, channels, kernel_size=3, stride=2, latent_dim=2):
super(ConvVAE,self).__init__()
self.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
self.latent_dim=latent_dim
last_channel = channels[0]
last_size = input_size
enc = []
#conv size calc
convcalc = lambda h,k,s,p: math.floor(((h+2*p-k)/s+1))
for c in channels[1:]:
enc.append(nn.Conv2d(last_channel,c,kernel_size,stride=stride,padding=1))
enc.append(nn.ReLU())
last_channel=c
last_size = convcalc(last_size,kernel_size,stride,1)
enc.append(nn.Flatten())
enc.append(nn.Linear(last_size*last_size*channels[-1],latent_dim*2))
self.seq_enc = nn.Sequential(*enc)
channels.reverse()
dec = [nn.Linear(latent_dim, last_size*last_size*channels[0]),
nn.ReLU(),
nn.Unflatten(1,(64,7,7)),
nn.ConvTranspose2d(channels[0],channels[0],kernel_size,stride),
nn.ReLU(),]
deconvcalc = lambda h,k,s,p : (h-1)*s-2*p+k
last_size = deconvcalc(last_size,kernel_size,stride,0)
last_channel = channels[0]
for c in channels[1:-1]:
dec.append(nn.ConvTranspose2d(last_channel,c,kernel_size,stride))
dec.append(nn.ReLU())
last_channel=c
last_size = deconvcalc(last_size,kernel_size,stride,0)
if(last_size>=input_size):
dec.append(nn.Conv2d(channels[-2], channels[-1], last_size-input_size+1, 1))
elif(last_size<input_size):
dec.append(nn.ConvTranspose2d(channels[-2], channels[-1], input_size-last_size+1, 1))
dec.append(nn.Sigmoid())
self.seq_dec = nn.Sequential(*dec)
def encoder(self,x):
x = self.seq_enc(x)
mean = x[:,0:self.latent_dim]
logvar = x[:,self.latent_dim:self.latent_dim*2]
return mean, logvar
def decoder(self,z):
xhat = self.seq_dec(z)
return xhat
def sample(self,mean,logvar):
#epsilon ~ N(O,I)
eps = torch.randn_like(mean).to(self.DEVICE)
return mean+torch.exp(0.5*logvar)*eps
def forward(self,x):
mean,logvar = self.encoder(x)
z = self.sample(mean,logvar)
xhat = self.decoder(z)
return xhat,mean,logvar Train
Prepare
Hyperparameter, 데이터셋, 모델, optimizer를 다음과 같이 로드한다.
Copy BATCH_SIZE=100
EPOCHS=50
Z_DIM = 2
transform = transforms.Compose([transforms.ToTensor()])
train_dataset= datasets.MNIST('./',train=True,transform=transform,download=True)
test_dataset = datasets.MNIST('./',train=False,transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=BATCH_SIZE,shuffle=True,pin_memory=True)
test_loader = torch.utils.data.DataLoader(test_dataset,batch_size=BATCH_SIZE,shuffle=False,pin_memory=True)
model = ConvVAE(input_size=28,channels=[1,32,64],latent_dim=Z_DIM)
optim = torch.optim.Adam(model.parameters(),lr=1e-3) Loss Function
Loss는 다음과 같이 정의한다. 미니배치의 크기가 100으로 충분히 크므로 L을 1로 해도 무관하다. 다음을 배치 하나 속의 데이터 포인트 100개에 대해 평균낸다(Monte Carlo Estimation).
-\tilde \mathcal L^M(\theta,\phi;\bold x^{(i)})=-\frac 1 2\sum^J_{j=1}\left(1 +\log((\sigma_j^{(i)})^2) -(\mu_j^{(i)})^2-(\sigma_j^{(i)})^2 \right)-\log p_{\bm \theta}(\bold x^{(i)}|\bold z^{(i,l)})
− log p θ ( x ( i ) ∣ z ( i , l ) ) -\log p_{\bm \theta}(\bold x^{(i)}|\bold z^{(i,l)}) − log p θ ( x ( i ) ∣ z ( i , l ) )
코드로 나타내면
Copy def lossfunc(mean,logvar,x,xhat):
BCE = torch.mean(torch.sum(F.binary_cross_entropy(xhat,x,reduction='none'),(1,2,3)))
KLD = torch.mean(-0.5*torch.sum(1+logvar-mean.pow(2)-logvar.exp(),1),0)
return BCE+KLD Reduction은 과정에 오류가 없도록 수동으로 해주었다.
Iteration
Train 과정의 코드는 기존것과 크게 다르지 않다.
Copy def train(epoch, model, optim, loader):
losses = 0
model.train()
for idx, (x,_) in enumerate(loader):
x = x.to(DEVICE)
optim.zero_grad()
xhat, mean,logvar = model(x)
loss = lossfunc(mean,logvar,x,xhat)
loss.backward()
optim.step()
losses+= loss.item()
if idx%200 ==0:
print(f'{100.*idx/len(loader):.1f}%, loss={loss.item():.2f}')
return losses/len(loader)
def test(epoch, model, loader):
losses = 0
model.eval()
for idx, (x,_) in enumerate(loader):
x = x.to(DEVICE)
xhat, mean,logvar = model(x)
with torch.no_grad():
loss = lossfunc(mean,logvar,x,xhat)
losses+= loss.item()
return losses/len(loader)
train_losses = []
test_losses = []
model.to(DEVICE)
for epoch in range(1,EPOCHS+1):
train_loss = train(epoch,model,optim,train_loader)
print(f'EPOCH {epoch}, train_loss={train_loss:.2f}',end='')
train_losses.append(train_loss)
test_loss = test(epoch,model,test_loader)
print(f', test_loss={test_loss:.2f}\n')
test_losses.append(test_loss)
if epoch%10 ==0:
torch.save(model.state_dict(),f'./checkpoint2/check{epoch}.pt')
print(train_losses)
print(test_losses) Result
학습 후, 잠재 변수와 MNIST 숫자 레이블이 잠재변수 공간 어디에 위치하는지 시각화했다. 학습이 잘 되었다면 올바르게 클러스터링되어야 한다.
Z_DIM=2
잠재변수가 두 개의 실수일 때 결과이다.
과적합 이전에 적절히 멈추었다.
시각화 결과,
점이 0 주변에 올바르게 분포해 했으나, 일부 숫자는 잘 분리되지 않았다.
100개를 Train Data에서 뽑아 적절히 생성했다. 결과, 다음과 같은 숫자가 생성되었다.
생성된 숫자가 애매하게 보이는 것이 많고, 압축 및 복원 과정에서 어려움이 있었다.
Z_DIM=4
잠재 변수가 4개의 실수일 때 결과이다.
Z_DIM=2에서보다 더 낮은 Loss를 보였다!
4차원이므로 시각화가 까다롭다. 따라서 T-SNE를 이용해 차원축소했다.
각 숫자가 더욱 잘 모여있는 것으로 보인다.
생성된 이미지들은 실제로 더 또렷한 숫자들로 보였으며, 0부터 9까지 각각 또렷하게 보이는 숫자들이 적어도 하나씩 존재했다. 더 잘 최적화되었다.
