구현해보기

PyTorch

AEVB / SGVB를 적용한 Variational Autoencoder를 구현한다.

Model

본래 VAE는 MLP로 구성되지만, 이 구현에선 컨볼루션 레이어를 사용한다. latent_dim은 잠재변수 z의 차원이다. 2와 4에서 실험할 것이다.

Encoder

Layer

Params

Conv+ReLU

1 -> 32, size=3x3, stride=2, pad=1

Conv+ReLU

32 -> 64, size=3x3, stride=2, pad=1

Flatten

Linear(mean)

3136 -> latent_dim

Linear(logvar)

3136 -> latent_dim

정규분포의 분산은 항상 양수인데, linear 레이어의 출력은 실수값이다. 실수 범위를 가지는 로그분산( $\log \sigma^2$ )을 출력한 후 $\sigma=\exp(\frac 1 2 \log\sigma^2)$ 을계산한다. 정규분포를 가정한 네트워크이다.

Sampling

torch.randn_like 로 epsilon 샘플 후, 이를 이용해 z를 계산한다.

Decoder

베르누이 분포를 가정한다. 출력이 0과 1 사이의 픽셀값이다.

Layer

Params

Linear+ReLU

latent_dim -> 3136

Reshape

(3136,) -> (64,7,7)

Deconv+ReLU

64 -> 64, size=3x3, stride=2

Deconv+ReLU

64 -> 32, size=3x3, stride=2

Conv

32 -> 1, size=4x4, stride=1

ConvTranspose2d로 차원(28x28)을 맟출 수 없었기에 최종적으로 Conv2d를 한번 더 해준다.

모델 클래스의 코드는 다음과 같다. 21/07/26 수정, 모델을 더 유연하게 변환했다.

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를 다음과 같이 로드한다.

Param

Value

Batch Size

100

Epochs

Optimizer / lr

Adam / 0.001

Z의 차원

2 또는 4

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)})

코드로 나타내면

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 과정의 코드는 기존것과 크게 다르지 않다.

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)