Hi every one,
I am writing a script for CNN VAE following PyTorch’s VAE tutorial.
But the error like the title makes me hard.
After searching, the corresponding error seems to be a problem between the gpu tensor and the cpu tensor.
I do not know which part of my script to fix below.
I’ve found that it works with cpu, but I’m confused about making the gpu tensor part.
============================================================================
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torchvision.utils import save_image
from torch.autograd import Variable
Device configuration
device = torch.device(‘cuda:0’)
#device =‘cpu’
Create a directory if not exists
sample_dir = ‘./samples’
if not os.path.exists(sample_dir):
os.makedirs(sample_dir)
Hyper-parameters
num_epochs = 15
batch_size = 64
learning_rate = 1e-3
IMAGE dataset
dataset = torchvision.datasets.ImageFolder(root="…/LipNet/mouth/cut/s3/",transform=transforms.Compose([transforms.Resize((64,64)),transforms.ToTensor()]))
#dataset = torchvision.datasets.ImageFolder(root="…/LipNet/mouth/cut/s3/",transform=transforms.ToTensor())
Data loader
data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size)
VAE model
class Flatten(nn.Module):
def forward(self, input):
return input.view(input.size(0), -1)
class UnFlatten(nn.Module):
def forward(self, input, size=1024):
return input.view(input.size(0), size, 1, 1)
class VAE(nn.Module):
def init(self, image_channels=3, h_dim=1024, z_dim=64):
super(VAE, self).init()
self.encoder = nn.Sequential(
nn.Conv2d(image_channels, 32, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(32, 64, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=4, stride=2),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=4, stride=2),
nn.ReLU(),
Flatten()
)
self.fc1 = nn.Linear(h_dim, z_dim)
self.fc2 = nn.Linear(h_dim, z_dim)
self.fc3 = nn.Linear(z_dim, h_dim)
self.decoder = nn.Sequential(
UnFlatten(),
nn.ConvTranspose2d(h_dim, 128, kernel_size=5, stride=2),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, kernel_size=5, stride=2),
nn.ReLU(),
nn.ConvTranspose2d(64, 32, kernel_size=6, stride=2),
nn.ReLU(),
nn.ConvTranspose2d(32, image_channels, kernel_size=6, stride=2),
nn.Sigmoid(),
)
def reparameterize(self, mu, logvar):
std = logvar.mul(0.5).exp_()
# return torch.normal(mu, std)
esp = torch.randn(*mu.size())
z = mu + std * esp
return z
def bottleneck(self, h):
mu, logvar = self.fc1(h), self.fc2(h)
z = self.reparameterize(mu, logvar)
return z, mu, logvar
def representation(self, x):
return self.bottleneck(self.encoder(x))[0]
def forward(self, x):
h = self.encoder(x)
z, mu, logvar = self.bottleneck(h)
z = self.fc3(z)
return self.decoder(z), mu, logvar
model = VAE().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
for epoch in range(num_epochs):
for i, (x, _) in enumerate(data_loader):
# Forward pass
print(“1”)
x = x.to(device)
print(“2”)
x_reconst, mu, log_var = model(x)
print(“3”)
reconst_loss = F.binary_cross_entropy(x_reconst, x, size_average=False)
print(“4”)
kl_div = - 0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
print(“5”)
# Backprop and optimize
loss = reconst_loss + kl_div
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 10 == 0:
print ("Epoch[{}/{}], step[{}/{}], Reconst Loss: {:.4f}, KL Div: {:.4f}"
.format(epoch+1, num_epochs, i+1, len(data_loader), reconst_loss.item(), kl_div.item()))
with torch.no_grad():
# Save the sampled images
z = torch.randn(batch_size, z_dim).to(device)
out = model.decode(z).view(-1, 1, 100, 50)
save_image(out, os.path.join(sample_dir, 'sampled-{}.png'.format(epoch+1)))
# Save the reconstructed images
out, _, _ = model(x)
x_concat = torch.cat([x.view(-1, 1, 100, 50), out.view(-1, 1, 100, 50)], dim=3)
save_image(x_concat, os.path.join(sample_dir, 'reconst-{}.png'.format(epoch+1)))
============================================================================