def train(self):
config = self.config
print(config.train_batch_size % len(self.unlabeled_loader))
self.param_init()
self.iter_cnt = 0
iter, min_dev_incorrect = 0, 1e6
monitor = OrderedDict()
batch_per_epoch = int((len(self.unlabeled_loader) + config.train_batch_size - 1) // config.train_batch_size)
while True:
if iter % batch_per_epoch == 0:
epoch = iter // batch_per_epoch
if epoch >= config.max_epochs:
break
epoch_ratio = float(epoch) / float(config.max_epochs)
# use another outer max to prevent any float computation precision problem
self.dis_optimizer.param_groups[0]['lr'] = config.dis_lr * max(0., min(3. * (1. - epoch_ratio), 1.))
self.gen_optimizer.param_groups[0]['lr'] = config.gen_lr * max(0., min(3. * (1. - epoch_ratio), 1.))
iter_vals = self._train(iter=iter)
for k, v in iter_vals.items():
if not (k) in monitor:
monitor[k] = 0.
monitor[k] += v
if iter % config.vis_period == 0:
self.visualize()
if iter % config.eval_period == 0:
train_loss, train_incorrect = self.eval(self.labeled_loader)
dev_loss, dev_incorrect = self.eval(self.dev_loader)
min_dev_incorrect = min(min_dev_incorrect, dev_incorrect)
disp_str = '#{}\ttrain: {:.4f}, {} | dev: {:.4f}, {} | best: {}'.format(
iter, train_loss, train_incorrect, dev_loss, dev_incorrect, min_dev_incorrect)
for k, v in monitor.items():
disp_str += ' | {}: {:.4f}'.format(k, v / config.eval_period)
disp_str += ' | lr: dis {:.5f}, gen {:.5f}'.format(
self.dis_optimizer.param_groups[0]['lr'], self.gen_optimizer.param_groups[0]['lr'])
monitor = OrderedDict()
print(disp_str)
noise = Variable(torch.Tensor(400, self.config.noise_size).uniform_().cuda(), volatile=True)
images = self.gen(noise)
images = (images - 0.5) // 0.5
images = images.view(-1, 1, 28, 28)
logits = self.pixelcnn(images)
log_probs = - pixelcnn_loss.discretized_mix_logistic_loss_c1(images.permute(0, 2, 3, 1), logits.permute(0, 2, 3, 1), sum_all=False).data.cpu()
gen_ploss_stats = log_probs.min(), log_probs.max(), log_probs.mean(), log_probs.var()
print ('gen stats', gen_ploss_stats)
print ('unl stats', self.unl_ploss_stats)
iter += 1
self.iter_cnt += 1
def __init__(self, config, args):
self.config = config
for k, v in list(args.__dict__.items()):
setattr(self.config, k, v)
setattr(self.config, 'save_dir', '{}_log'.format(self.config.dataset))
self.dis = mnist_model.Discriminative(
noise_size=config.noise_size, num_label=config.num_label).cuda()
self.gen = mnist_model.Generator(image_size=config.image_size,
noise_size=config.noise_size).cuda()
self.dis_optimizer = optim.Adam(self.dis.parameters(),
lr=config.dis_lr,
betas=(0.5, 0.9999))
self.gen_optimizer = optim.Adam(self.gen.parameters(),
lr=config.gen_lr,
betas=(0.0, 0.9999))
self.pixelcnn = pixelcnn.PixelCNN(nr_resnet=3,
disable_third=True,
dropout_p=0.0,
n_channel=1,
image_wh=28).cuda()
self.pixelcnn.load_state_dict(torch.load(config.pixelcnn_path))
self.d_criterion = nn.CrossEntropyLoss()
self.labeled_loader, self.unlabeled_loader, self.unlabeled_loader2, self.dev_loader, self.special_set = data.get_mnist_loaders(
config)
if not os.path.exists(self.config.save_dir):
os.makedirs(self.config.save_dir)
log_probs_list = []
for dev_images, _ in self.dev_loader.get_iter():
dev_images = Variable(dev_images.cuda(), requires_grad=False)
dev_images = (dev_images - 0.5) / 0.5
dev_images = dev_images.view(-1, 1, 28, 28)
with torch.no_grad():
logits = self.pixelcnn(dev_images)
log_probs = -pixelcnn_loss.discretized_mix_logistic_loss_c1(
dev_images.permute(0, 2, 3, 1),
logits.permute(0, 2, 3, 1),
sum_all=False)
log_probs = log_probs.data.cpu().numpy()
log_probs_list.append(log_probs)
log_probs = np.concatenate(log_probs_list, axis=0)
self.unl_ploss_stats = log_probs.min(), log_probs.max(
), log_probs.mean(), log_probs.var()
cut_point = int(log_probs.shape[0] * 0.1)
self.ploss_th = float(np.partition(log_probs, cut_point)[cut_point])
print('ploss_th', self.ploss_th)
print(self.dis)
def _train(self, labeled=None, vis=False, iter=0):
config = self.config
self.dis.train()
self.gen.train()
lab_images, lab_labels = next(self.labeled_loader)
lab_images, lab_labels = Variable(lab_images.cuda()), Variable(
lab_labels.cuda())
unl_images, _ = next(self.unlabeled_loader)
unl_images = Variable(unl_images.cuda())
noise = Variable(
torch.Tensor(unl_images.size(0),
config.noise_size).uniform_().cuda())
gen_images = self.gen(noise)
lab_logits = self.dis(lab_images)
unl_logits = self.dis(unl_images)
gen_logits = self.dis(gen_images.detach())
# Standard classification loss
lab_loss = self.d_criterion(lab_logits, lab_labels)
unl_logsumexp = log_sum_exp(unl_logits)
gen_logsumexp = log_sum_exp(gen_logits)
unl_acc = torch.mean(
nn.functional.sigmoid(unl_logsumexp.detach()).gt(0.5).float())
gen_acc = torch.mean(
nn.functional.sigmoid(gen_logsumexp.detach()).gt(0.5).float())
# This is the typical GAN cost, where sumexp(logits) is seen as the input to the sigmoid
true_loss = -0.5 * torch.mean(unl_logsumexp) + 0.5 * torch.mean(
F.softplus(unl_logsumexp))
fake_loss = 0.5 * torch.mean(F.softplus(gen_logsumexp))
unl_loss = true_loss + fake_loss
d_loss = lab_loss + unl_loss
self.dis_optimizer.zero_grad()
d_loss.backward()
self.dis_optimizer.step()
##### train Gen and Enc
unl_images, _ = next(self.unlabeled_loader2)
unl_images = Variable(unl_images.cuda())
noise = Variable(
torch.Tensor(unl_images.size(0),
config.noise_size).uniform_().cuda())
gen_images = self.gen(noise)
unl_feat = self.dis(unl_images, feat=True)
gen_feat = self.dis(gen_images, feat=True)
fm_loss = torch.mean(
(torch.mean(gen_feat, 0) - torch.mean(unl_feat, 0))**2)
gen_loss = fm_loss
if iter > 9000 and random.random() < config.p_loss_prob:
noise = Variable(
torch.Tensor(30, config.noise_size).uniform_().cuda())
gen_images = self.gen(noise)
gen_images = (gen_images - 0.5) / 0.5
gen_images = gen_images.view(-1, 1, 28, 28)
logits = self.pixelcnn(gen_images)
log_probs = -pixelcnn_loss.discretized_mix_logistic_loss_c1(
gen_images.permute(0, 2, 3, 1),
logits.permute(0, 2, 3, 1),
sum_all=False)
p_loss = torch.max(
log_probs - self.ploss_th,
Variable(torch.cuda.FloatTensor(log_probs.size()).fill_(0.0)))
non_zero_cnt = float((p_loss > 0).sum().data.cpu()[0])
if non_zero_cnt > 0:
p_loss = p_loss.sum() / non_zero_cnt * config.p_loss_weight
else:
p_loss = 0
else:
p_loss = 0
loss = gen_loss + p_loss
self.gen_optimizer.zero_grad()
loss.backward()
self.gen_optimizer.step()
monitor_dict = OrderedDict([
('unl accuracy', unl_acc.data[0]),
('gen accuracy', gen_acc.data[0]), ('lab loss', lab_loss.data[0]),
('unl loss', unl_loss.data[0]), ('true loss', true_loss.data[0]),
('fake loss', fake_loss.data[0]), ('gen loss', gen_loss.data[0]),
('p loss', p_loss.data[0] if hasattr(p_loss, 'data') else 0.0)
])
return monitor_dict