def __init__(self, in_dim, encoder_width, decoder_width, latent_dim, device=torch.device('cpu')):
# device
self.name = 'VAE'
self.device = device
self.latent_dim = latent_dim
# initialize encoder/decoder weights and biases
self.weights, self.biases = self.init_vae_params(in_dim, encoder_width, decoder_width, latent_dim)
# config dataset
mnist = MNIST()
self.train_data = mnist.get(train=True)
self.test_data = mnist.get(train=False)
import torch
from data.dataset import MNIST
data_loader_train = torch.utils.data.DataLoader(MNIST(),
batch_size=2375,
shuffle=False)
def train(args, config):
vis = Visualizer()
train_set = MNIST(data_path=config.train_data_path,
label_path=config.train_label_path,
config=config,
mode='train')
valid_set = MNIST(data_path=config.train_data_path,
label_path=config.train_label_path,
config=config,
mode='valid')
train_dataloader = DataLoader(train_set,
config.batch_size,
shuffle=True,
num_workers=config.num_workers)
valid_dataloader = DataLoader(valid_set,
config.batch_size,
shuffle=False,
num_workers=config.num_workers)
model = getattr(network, args.model)().eval()
if args.load_model_path:
model.load(args.load_model_path)
if args.use_gpu:
model.cuda()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=config.lr)
train_loss_meter, valid_loss_meter = meter.AverageValueMeter(
), meter.AverageValueMeter()
train_confusion_matrix, valid_confusion_matrix = meter.ConfusionMeter(
10), meter.ConfusionMeter(10)
best_valid_loss = 1e5
best_epoch = 0
dist_to_best = 0
time_begin = time.clock()
for epoch in range(config.epoch):
# train
model.train()
train_loss_meter.reset()
train_confusion_matrix.reset()
for _iter, (train_data, train_target) in enumerate(train_dataloader):
if args.use_gpu:
train_data = train_data.cuda()
train_target = train_target.cuda()
optimizer.zero_grad()
train_logits, train_output = model(train_data)
train_loss = criterion(train_logits, train_target)
train_loss.backward()
optimizer.step()
train_loss_meter.add(train_loss.item())
train_confusion_matrix.add(train_logits.data, train_target.data)
if _iter % config.print_freq == 0:
vis.plot('train_loss', train_loss_meter.value()[0])
model.save(path=os.path.join(args.ckpts_dir, 'model_{0}.pth'.format(
str(epoch))))
# valid
model.eval()
valid_loss_meter.reset()
valid_confusion_matrix.reset()
for _iter, (valid_data, valid_target) in enumerate(valid_dataloader):
if args.use_gpu:
valid_data = valid_data.cuda()
valid_target = valid_target.cuda()
valid_logits, valid_output = model(valid_data)
valid_loss = criterion(valid_logits, valid_target)
valid_loss_meter.add(valid_loss.item())
valid_confusion_matrix.add(valid_logits.detach().squeeze(),
valid_target.type(t.LongTensor))
valid_cm = valid_confusion_matrix.value()
valid_accuracy = 100. * (valid_cm.diagonal().sum()) / (valid_cm.sum())
vis.plot('valid_accuracy', valid_accuracy)
vis.log(
"epoch:{epoch}, train_loss:{train_loss}, train_cm:{train_cm}, valid_loss:{valid_loss}, valid_cm:{valid_cm}, valid_accuracy:{valid_accuracy}"
.format(epoch=epoch,
train_loss=train_loss_meter.value()[0],
train_cm=str(train_confusion_matrix.value()),
valid_loss=valid_loss_meter.value()[0],
valid_cm=str(valid_cm),
valid_accuracy=valid_accuracy))
print(
"epoch:{epoch}, train_loss:{train_loss}, valid_loss:{valid_loss}, valid_accuracy:{valid_accuracy}"
.format(epoch=epoch,
train_loss=train_loss_meter.value()[0],
valid_loss=valid_loss_meter.value()[0],
valid_accuracy=valid_accuracy))
print("train_cm:\n{train_cm}\n\nvalid_cm:\n{valid_cm}".format(
train_cm=str(train_confusion_matrix.value()),
valid_cm=str(valid_cm),
))
# early stop
if valid_loss_meter.value()[0] < best_valid_loss:
best_epoch = epoch
best_valid_loss = valid_loss_meter.value()[0]
dist_to_best = 0
dist_to_best += 1
if dist_to_best > 4:
break
model.save(path=os.path.join(args.ckpts_dir, 'model.pth'))
vis.save()
print("save model successfully")
print("best epoch: ", best_epoch)
print("best valid loss: ", best_valid_loss)
time_end = time.clock()
print('time cost: %.2f' % (time_end - time_begin))
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from model import model_dcgan
import torch.utils.data as data
import numpy as np
from data.dataset import MNIST
import torchvision.utils as vutils
batch_size = 50
mnist = MNIST()
dataloader = data.DataLoader(mnist, batch_size=batch_size, shuffle=True)
gnet = model_dcgan.generator()
dnet = model_dcgan.discriminator()
if torch.cuda.is_available():
gnet = gnet.cuda()
dnet = dnet.cuda()
g_optim = optim.Adam(gnet.parameters(), lr=0.0002, betas=(0.5, 0.99))
d_optim = optim.Adam(dnet.parameters(), lr=0.0002, betas=(0.5, 0.99))
crit = nn.BCELoss()
if torch.cuda.is_available():
crit = crit.cuda()
fixed_z = Variable(torch.randn(100, 100, 1, 1))
fixed_y = Variable(