Hi everyone. I am trying to create a VAE which takes in different images from a dataset as the input, but am getting an error.
Here is the code:
from skimage.io import imread
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import numpy as np
import torch
import torch.utils.data
from torch import nn, optim
from torch.autograd import Variable
from torch.nn import functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
from PIL import Image
import os
import ast
import numpy as np
import pywt
import matplotlib.pyplot as plt
from PIL import Image as PImage
os.chdir(‘D:/bob/’)
SAVEDIR = “D:/bob/myfiles/image/output/”
IMAGEDIR=“D:/bob/myfiles/dataset”
device=“cpu”
class VAE(nn.Module):
def init(self, x_dim, h_dim, z_dim):
super(VAE, self).init()
# Encoder
self.fc1 = nn.Linear(x_dim, h_dim)
self.fc2_mu = nn.Linear(h_dim, z_dim)
self.fc2_logvar = nn.Linear(h_dim, z_dim)
# Decoder
self.fc3 = nn.Linear(z_dim, h_dim)
self.fc4 = nn.Linear(h_dim, x_dim)
def encoder(self, x):
h = F.relu(self.fc1(x))
return self.fc2_mu(h), self.fc2_logvar(h)
def sampling(self, mu, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return eps.mul(std).add_(mu) # return z sample
def decoder(self, z):
h = F.relu(self.fc3(z))
return F.sigmoid(self.fc4(h))
def forward(self, x):
mu, logvar = self.encoder(x.view(-1, 4096))
z = self.sampling(mu, logvar)
return self.decoder(z), mu,
vae = VAE(x_dim=4096, h_dim=400, z_dim=20).to(device)
optimizer = optim.Adam(vae.parameters(), lr=0.001)
reconstruction error + KL divergence losses
def loss_function(recon_x, x, mu, logvar):
BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction=‘sum’)
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + KLD
dirpath = os.listdir(IMAGEDIR)
def train(epoch):
vae.train()
train_loss = 0
for subdir in dirpath:
full_subdir= IMAGEDIR + ‘/’ + subdir + ‘/’
imageList = os.listdir(full_subdir)
for image in imageList:
fullpath = full_subdir + image
inputImage = imread(fullpath)[:,:]
optimizer.zero_grad()
recon_batch, mu, logvar = vae(inputImage)
loss = loss_function(recon_batch, inputImage, mu, logvar)
loss.backward()
train_loss += loss.item()
optimizer.step()
train(7)
When I try to run the code, this is the error I get:
ValueError: Type must be a sub-type of ndarray type
This comes from the line that says: mu, logvar = self.encoder(x.view(-1, 4096)) in the forward function. Does anyone have any solutions as to how I can resolve this issue so my code will run properly? Thanks!