MNIST データセットを使用して DCGAN モデルをトレーニングしようとしていましたが、gen.state_dict()トレーニングが終了した後にロードできません。
import torch
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
import torchvision
import os
from torch.autograd import Variable
workspace_dir = '/content/drive/My Drive/practice'
device=torch.device('cuda' if torch.cuda.is_available else 'cpu')
print(device)
img_size=64
channel_img=1
lr=2e-4
batch_size=128
z_dim=100
epochs=10
features_gen=64
features_disc=64
save_dir = os.path.join(workspace_dir, 'logs')
os.makedirs(save_dir, exist_ok=True)
import matplotlib.pyplot as plt
transforms=transforms.Compose([transforms.Resize(img_size),transforms.ToTensor(),transforms.Normalize(mean=(0.5,),std=(0.5,))])
train_data=datasets.MNIST(root='dataset/',train=True,transform=transforms,download=True)
train_loader=torch.utils.data.DataLoader(train_data,batch_size=batch_size,shuffle=True)
count=0
for x,y in train_loader:
if count==5:
break
print(x.shape,y.shape)
count+=1
class Discriminator(nn.Module):
def __init__(self,channels_img,features_d):
super(Discriminator,self).__init__()
self.disc=nn.Sequential(
#input:N * channels_img * 64 *64
nn.Conv2d(channels_img,features_d,4,2,1),#paper didn't use batchnorm in the early layers in the discriminator features_d* 32 *32
nn.LeakyReLU(0.2),
self._block(features_d,features_d*2,4,2,1),#features_d*2 *16 *16
self._block(features_d*2,features_d*4,4,2,1),#features_d*4 *8 *8
self._block(features_d*4,features_d*8,4,2,1), #features_d*8 *4 *4
nn.Conv2d(features_d*8,1,4,2,0),#1 * 1 *1
nn.Sigmoid()
)
def _block(self,in_channels,out_channels,kernel_size,stride,padding):
return nn.Sequential(
nn.Conv2d(in_channels,out_channels,kernel_size,stride,padding,bias=False),
nn.BatchNorm2d(out_channels),
nn.LeakyReLU(0.2)
)
def forward(self,x):
return self.disc(x)
class Generator(nn.Module):
def __init__(self,Z_dim,channels_img,features_g):
super(Generator,self).__init__()
self.gen=nn.Sequential(
#input :n * z_dim * 1 *1
self._block(Z_dim,features_g*16,4,1,0),#features_g*16 * 4 * 4
self._block(features_g*16,features_g*8,4,2,1),#features_g*8 * 8 * 8
self._block(features_g*8,features_g*4,4,2,1),#features_g*4 * 16 * 16
self._block(features_g*4,features_g*2,4,2,1),#features_g*2 * 32 * 32
nn.ConvTranspose2d(features_g*2,channels_img,4,2,1), #
nn.Tanh()# [-1 to 1] normalize the image
)
def _block(self,in_channels,out_channels,kernel_size,stride,padding):
return nn.Sequential(
nn.ConvTranspose2d(in_channels,out_channels,kernel_size,stride,padding,bias=False),#w'=(w-1)*s-2p+k
nn.BatchNorm2d(out_channels),
nn.ReLU()
)
def forward(self,x):
return self.gen(x)
def initialize_weights(model):
for m in model.modules():
if isinstance(m,(nn.Conv2d,nn.ConvTranspose2d,nn.BatchNorm2d)):
nn.init.normal_(m.weight.data,0.0,0.02)
gen=Generator(z_dim,channel_img,features_gen).to(device)
disc=Discriminator(channel_img,features_disc).to(device)
initialize_weights(gen)
initialize_weights(disc)
opt_gen=torch.optim.Adam(gen.parameters(),lr=lr,betas=(0.5,0.999))
opt_disc=torch.optim.Adam(disc.parameters(),lr=lr,betas=(0.5,0.999))
criterion=nn.BCELoss()
#fixed_noise=torch.randn(32,z_dim,1,1).to(device)
#writer_real=SummaryWriter(f"logs/real")
#writer_fake=SummaryWriter(f"logs/fake")
step=0
gen.train()
disc.train()
z_sample = Variable(torch.randn(100, z_dim,1,1)).cuda()
for epoch in range(2):
for batch_idx,(real,_) in enumerate(train_loader):
real=real.to(device)
noise=torch.randn((batch_size,z_dim,1,1)).to(device)
fake=gen(noise)
#Train Discriminator max log(D(x)) + log(1-D(G(z)))
disc_real=disc(real).reshape(-1)
loss_disc_real=criterion(disc_real,torch.ones_like(disc_real))
disc_fake=disc(fake).reshape(-1)
loss_disc_fake=criterion(disc_fake,torch.zeros_like(disc_fake))
loss_disc=(loss_disc_fake+loss_disc_real)/2
disc.zero_grad()
loss_disc.backward(retain_graph=True)
opt_disc.step()
#Train Generator min log(1-D(G(z))) <--> max log(D(G(z)))
output=disc(fake).reshape(-1)
loss_gen=criterion(output,torch.ones_like(output))
gen.zero_grad()
loss_gen.backward()
opt_gen.step()
print(f'\rEpoch [{epoch+1}/{3}] {batch_idx+1}/{len(train_loader)} Loss_D: {loss_disc.item():.4f} Loss_G: {loss_gen.item():.4f}', end='')
gen.eval()
f_imgs_sample = (gen(z_sample).data + 1) / 2.0
filename = os.path.join(save_dir, f'Epoch_{epoch+1:03d}.jpg')
torchvision.utils.save_image(f_imgs_sample, filename, nrow=10)
print(f' | Save some samples to {filename}.')
# show generated image
grid_img = torchvision.utils.make_grid(f_imgs_sample.cpu(), nrow=10)
plt.figure(figsize=(10,10))
plt.imshow(grid_img.permute(1, 2, 0))
plt.show()
gen.train()
torch.save(gen.state_dict(), os.path.join(workspace_dir, f'dcgan_d.pth'))
torch.save(disc.state_dict(), os.path.join(workspace_dir, f'dcgan_g.pth'))
このステップで gen state_dict を読み込めません:
# load pretrained model
#gen = Generator(z_dim,1,64)
gen=Generator(z_dim,channel_img,features_gen).to(device)
gen.load_state_dict(torch.load('/content/drive/My Drive/practice/dcgan_g.pth'))
gen.eval()
gen.cuda()
エラーは次のとおりです。
RuntimeError Traceback (most recent call last)
<ipython-input-18-4bda27faa444> in <module>()
5
6 #gen.load_state_dict(torch.load('/content/drive/My Drive/practice/dcgan_g.pth'))
----> 7 gen.load_state_dict(torch.load(os.path.join(workspace_dir, 'dcgan_g.pth')))
8 #/content/drive/My Drive/practice/dcgan_g.pth
9 gen.eval()
/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
1050 if len(error_msgs) > 0:
1051 raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-> 1052 self.__class__.__name__, "\n\t".join(error_msgs)))
1053 return _IncompatibleKeys(missing_keys, unexpected_keys)
1054
***RuntimeError: Error(s) in loading state_dict for Generator:
Missing key(s) in state_dict***: "gen.0.0.weight", "gen.0.1.weight", "gen.0.1.bias", "gen.0.1.running_mean", "gen.0.1.running_var", "gen.1.0.weight", "gen.1.1.weight", "gen.1.1.bias", "gen.1.1.running_mean", "gen.1.1.running_var", "gen.2.0.weight", "gen.2.1.weight", "gen.2.1.bias", "gen.2.1.running_mean", "gen.2.1.running_var", "gen.3.0.weight", "gen.3.1.weight", "gen.3.1.bias", "gen.3.1.running_mean", "gen.3.1.running_var", "gen.4.weight", "gen.4.bias".
Unexpected key(s) in state_dict: "disc.0.weight", "disc.0.bias", "disc.2.0.weight", "disc.2.1.weight", "disc.2.1.bias", "disc.2.1.running_mean", "disc.2.1.running_var", "disc.2.1.num_batches_tracked", "disc.3.0.weight", "disc.3.1.weight", "disc.3.1.bias", "disc.3.1.running_mean", "disc.3.1.running_var", "disc.3.1.num_batches_tracked", "disc.4.0.weight", "disc.4.1.weight", "disc.4.1.bias", "disc.4.1.running_mean", "disc.4.1.running_var", "disc.4.1.num_batches_tracked", "disc.5.weight", "disc.5.bias".