I met a problem with data type when transforming my own image training data to a torch.utils.data.Dataset
The Dataset is defined as belows
class FaceDataset(torch.utils.data.Dataset):
"""
读取本地动漫人脸数据
"""
def __init__(self):
data = np.loadtxt("Data1.csv",delimiter=',',skiprows=1,dtype='float32')
self.torch_data = torch.from_numpy(data)
self.len = data.shape[0]
def __getitem__(self, index):
return self.torch_data[index],1
def __len__(self):
return self.len
and the model
class CNN_Discriminator(nn.Module):
def __init__(self):
super(CNN_Discriminator, self).__init__()
self.conv1 = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=32, kernel_size=5, padding=2), # batch, 32, 96，96,
nn.LeakyReLU(0.2, True),
nn.AvgPool2d(2, stride=2), # batch, 32, 48, 48
)
self.conv2 = nn.Sequential(
nn.Conv2d(32, 64, 5, padding=2), # batch, 64, 48, 48
nn.LeakyReLU(0.2, True),
nn.AvgPool2d(2, stride=3) # batch, 64, 16, 16
)
self.fc = nn.Sequential(
nn.Linear(64 * 16 * 16, 1024),
nn.LeakyReLU(0.2, True),
nn.Linear(1024, 1),
nn.Sigmoid()
)
def forward(self, x):
'''
x: batch, width, height, channel=3
'''
x = self.conv1(x) # !!!! Code Went wrong Here!!
x = self.conv2(x)
x = x.view(x.size(0), -1)
x = self.fc(x)
return x
“main.py” if needed
import torch
import numpy as np
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms
from torchvision.utils import save_image
from torch.autograd import Variable
import os
import model
def to_img(x):
out = 0.5 * (x + 1)
out = out.clamp(0, 1)
out = out.view(-1, 3, 96, 96)
return out
class FaceDataset(torch.utils.data.Dataset):
"""
load imgs
"""
def __init__(self):
data = np.loadtxt("Data1.csv",delimiter=',',skiprows=1,dtype='float32')
self.torch_data = torch.from_numpy(data)
self.len = data.shape[0]
def __getitem__(self, index):
return self.torch_data[index],1
def __len__(self):
return self.len
batch_size = 2
num_epoch = 10
z_dimension = 512
dataSet = FaceDataset()
img_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
dataloader = torch.utils.data.DataLoader(
dataset=dataSet, batch_size=batch_size, shuffle=True,num_workers=8)
G = model.CNN_Generator(z_dimension,15*192*192)
D = model.CNN_Discriminator()
if torch.cuda.is_available():
D = D.cuda()
G = G.cuda()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Binary cross entropy loss and optimizer
criterion = nn.BCELoss()
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0002)
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0002)
for epoch in range(num_epoch):
for i, (img,_) in enumerate(dataloader):
num_img = img.size(0)
#train discriminator
img = img.view(num_img,3,96,96)
img = img.double()
real_img = Variable(img).to(device)
real_label = Variable(torch.ones(num_img)).to(device)
fake_label = Variable(torch.zeros(num_img)).to(device)
# compute loss of real_img
real_out = D(real_img) # !!!!!!Problem Here
d_loss_real = criterion(real_out, real_label)
real_scores = real_out # closer to 1 means better
# compute loss of fake_img
z = Variable(torch.randn(num_img, z_dimension))#.cuda()
fake_img = G(z)
fake_out = D(fake_img)
d_loss_fake = criterion(fake_out, fake_label)
fake_scores = fake_out # closer to 0 means better
# bp and optimize
d_loss = d_loss_real + d_loss_fake
d_optimizer.zero_grad()
d_loss.backward()
d_optimizer.step()
# ===============train generator
# compute loss of fake_img
z = Variable(torch.randn(num_img, z_dimension)).to(device)
fake_img = G(z)
output = D(fake_img)
g_loss = criterion(output, real_label)
# bp and optimize
g_optimizer.zero_grad()
g_loss.backward()
g_optimizer.step()
print('Epoch [{}/{}], Batch {}/{},d_loss: {:.6f}, g_loss: {:.6f} '
'D real: {:.6f}, D fake: {:.6f}'.format(
epoch, num_epoch,i,dataSet.len/batch_size ,d_loss.data, g_loss.data,
real_scores.data.mean(), fake_scores.data.mean()))
if epoch == 0:
real_images = to_img(real_img.cpu().data)
save_image(real_images, './img/real_images.png')
fake_images = to_img(fake_img.cpu().data)
save_image(fake_images, './img/fake_images-{}.png'.format(epoch + 1))
torch.save(G.state_dict(), './generator.pth')
torch.save(D.state_dict(), './discriminator.pth')