def fetch_multimnist_image(_label):
if _label == EMPTY:
_label = ''
loader = torch.utils.data.DataLoader(datasets.MultiMNIST(
'./data',
train=False,
download=True,
transform=transforms.ToTensor(),
target_transform=charlist_tensor),
batch_size=1,
shuffle=True)
images = []
for image, label in loader:
if tensor_to_string(label.squeeze(0)) == _label:
images.append(image)
if len(images) == 0:
sys.exit('No images with label (%s) found.' % _label)
images = torch.cat(images).cpu().numpy()
ix = np.random.choice(np.arange(images.shape[0]))
image = images[ix]
image = torch.from_numpy(image).float()
image = image.unsqueeze(0)
return Variable(image, volatile=True)
type=str,
help='path to output directory of weak.py')
parser.add_argument('--cuda',
action='store_true',
default=False,
help='enables CUDA training')
args = parser.parse_args()
args.cuda = args.cuda and torch.cuda.is_available()
x, y1, y2 = [], [], []
for dir_path in glob(os.path.join(args.models_dir, '*')):
weak_perc = float(os.path.basename(dir_path).split('_')[-1])
loader = torch.utils.data.DataLoader(datasets.MultiMNIST(
'./data',
train=False,
download=True,
transform=transforms.ToTensor(),
target_transform=charlist_tensor),
batch_size=128,
shuffle=True)
vae = load_checkpoint(os.path.join(dir_path, 'model_best.pth.tar'),
use_cuda=args.cuda)
vae.eval()
weak_char_acc, weak_len_acc = test_multimnist(vae,
loader,
use_cuda=args.cuda,
verbose=False)
x.append(weak_perc)
y1.append(weak_char_acc)
y2.append(weak_len_acc)
print('Got accuracies for %s.' % dir_path)
def train_pipeline(out_dir,
weak_perc_m1,
weak_perc_m2,
n_latents=20,
batch_size=128,
epochs=20,
lr=1e-3,
log_interval=10,
cuda=False):
"""Pipeline to train and test MultimodalVAE on MNIST dataset. This is
identical to the code in train.py.
:param out_dir: directory to store trained models
:param weak_perc_m1: percent of time to show first modality
:param weak_perc_m2: percent of time to show second modality
:param n_latents: size of latent variable (default: 20)
:param batch_size: number of examples to show at once (default: 128)
:param epochs: number of loops over dataset (default: 20)
:param lr: learning rate (default: 1e-3)
:param log_interval: interval of printing (default: 10)
:param cuda: whether to use cuda or not (default: False)
"""
# create loaders for MNIST
train_loader = torch.utils.data.DataLoader(datasets.MultiMNIST(
'./data',
train=True,
download=True,
transform=transforms.ToTensor(),
target_transform=charlist_tensor),
batch_size=batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.MultiMNIST(
'./data',
train=False,
download=True,
transform=transforms.ToTensor(),
target_transform=charlist_tensor),
batch_size=batch_size,
shuffle=True)
# load multimodal VAE
vae = MultimodalVAE(n_latents=n_latents, use_cuda=cuda)
if cuda:
vae.cuda()
optimizer = optim.Adam(vae.parameters(), lr=lr)
def train(epoch):
random.seed(42)
np.random.seed(42) # important to have the same seed
# in order to make the same choices for weak supervision
# otherwise, we end up showing different examples over epochs
vae.train()
joint_loss_meter = AverageMeter()
image_loss_meter = AverageMeter()
text_loss_meter = AverageMeter()
for batch_idx, (image, text) in enumerate(train_loader):
if cuda:
image, text = image.cuda(), text.cuda()
image, text = Variable(image), Variable(text)
optimizer.zero_grad()
recon_image_1, recon_text_1, mu_1, logvar_1 = vae(image, text)
loss = loss_function(mu_1,
logvar_1,
recon_image=recon_image_1,
image=image,
recon_text=recon_text_1,
text=text,
kl_lambda=kl_lambda,
lambda_xy=1.,
lambda_yx=1.)
joint_loss_meter.update(loss.data[0], len(image))
# depending on this flip, we decide whether or not to show a modality
# versus another one.
flip = np.random.random()
if flip < weak_perc_m1:
recon_image_2, recon_text_2, mu_2, logvar_2 = vae(image=image)
loss_2 = loss_function(mu_2,
logvar_2,
recon_image=recon_image_2,
image=image,
recon_text=recon_text_2,
text=text,
kl_lambda=kl_lambda,
lambda_xy=1.,
lambda_yx=1.)
image_loss_meter.update(loss_2.data[0], len(image))
loss += loss_2
flip = np.random.random()
if flip < weak_perc_m2:
recon_image_3, recon_text_3, mu_3, logvar_3 = vae(text=text)
loss_3 = loss_function(mu_3,
logvar_3,
recon_image=recon_image_3,
image=image,
recon_text=recon_text_3,
text=text,
kl_lambda=kl_lambda,
lambda_xy=0.,
lambda_yx=1.)
text_loss_meter.update(loss_3.data[0], len(text))
loss += loss_3
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print(
'[Weak (Image) {:.0f}% | Weak (Text) {:.0f}%] Train Epoch: {} [{}/{} ({:.0f}%)]\tJoint Loss: {:.6f}\tImage Loss: {:.6f}\tText Loss: {:.6f}'
.format(100. * weak_perc_m1, 100. * weak_perc_m2, epoch,
batch_idx * len(image), len(train_loader.dataset),
100. * batch_idx / len(train_loader),
joint_loss_meter.avg, image_loss_meter.avg,
text_loss_meter.avg))
print(
'====> [Weak (Image) {:.0f}% | Weak (Text) {:.0f}%] Epoch: {} Joint loss: {:.4f}\tImage loss: {:.4f}\tText loss: {:.4f}'
.format(100. * weak_perc_m1, 100. * weak_perc_m2, epoch,
joint_loss_meter.avg, image_loss_meter.avg,
text_loss_meter.avg))
def test():
vae.eval()
test_joint_loss = 0
test_image_loss = 0
test_text_loss = 0
for batch_idx, (image, text) in enumerate(test_loader):
if cuda:
image, text = image.cuda(), text.cuda()
image, text = Variable(image), Variable(text)
# in test i always care about the joint loss -- so we don't anneal
# back joint examples as we do in train
recon_image_1, recon_text_1, mu_1, logvar_1 = vae(image, text)
recon_image_2, recon_text_2, mu_2, logvar_2 = vae(image=image)
recon_image_3, recon_text_3, mu_3, logvar_3 = vae(text=text)
loss_1 = loss_function(mu_1,
logvar_1,
recon_image=recon_image_1,
image=image,
recon_text=recon_text_1,
text=text,
kl_lambda=kl_lambda,
lambda_xy=1.,
lambda_yx=1.)
loss_2 = loss_function(mu_2,
logvar_2,
recon_image=recon_image_2,
image=image,
recon_text=recon_text_2,
text=text,
kl_lambda=kl_lambda,
lambda_xy=1.,
lambda_yx=1.)
loss_3 = loss_function(mu_3,
logvar_3,
recon_image=recon_image_3,
image=image,
recon_text=recon_text_3,
text=text,
kl_lambda=kl_lambda,
lambda_xy=0.,
lambda_yx=1.)
test_joint_loss += loss_1.data[0]
test_image_loss += loss_2.data[0]
test_text_loss += loss_3.data[0]
test_loss = test_joint_loss + test_image_loss + test_text_loss
test_joint_loss /= len(test_loader)
test_image_loss /= len(test_loader)
test_text_loss /= len(test_loader)
test_loss /= len(test_loader)
print(
'====> [Weak (Image) {:.0f}% | Weak (Text) {:.0f}%] Test joint loss: {:.4f}\timage loss: {:.4f}\ttext loss:{:.4f}'
.format(100. * weak_perc_m1, 100. * weak_perc_m2, test_joint_loss,
test_image_loss, test_text_loss))
return test_loss, (test_joint_loss, test_image_loss, test_text_loss)
best_loss = sys.maxint
schedule = iter([5e-5, 1e-4, 5e-4, 1e-3])
for epoch in range(1, epochs + 1):
if (epoch - 1) % 10 == 0:
try:
kl_lambda = next(schedule)
except:
pass
train(epoch)
loss, (joint_loss, image_loss, text_loss) = test()
is_best = loss < best_loss
best_loss = min(loss, best_loss)
save_checkpoint(
{
'state_dict': vae.state_dict(),
'best_loss': best_loss,
'joint_loss': joint_loss,
'image_loss': image_loss,
'text_loss': text_loss,
'optimizer': optimizer.state_dict(),
},
is_best,
folder=out_dir)
if not os.path.isdir('./trained_models'):
os.makedirs('./trained_models')
if not os.path.isdir('./trained_models/text_only'):
os.makedirs('./trained_models/text_only')
if not os.path.isdir('./results'):
os.makedirs('./results')
if not os.path.isdir('./results/text_only'):
os.makedirs('./results/text_only')
train_loader = torch.utils.data.DataLoader(
datasets.MultiMNIST('./data', train=True, download=True,
transform=transforms.ToTensor(),
target_transform=charlist_tensor),
batch_size=args.batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MultiMNIST('./data', train=False, download=True,
transform=transforms.ToTensor(),
target_transform=charlist_tensor),
batch_size=args.batch_size, shuffle=True)
vae = TextVAE(args.n_latents, use_cuda=args.cuda)
if args.cuda:
vae.cuda()
optimizer = optim.Adam(vae.parameters(), lr=args.lr)
args = parser.parse_args()
args.cuda = args.cuda and torch.cuda.is_available()
if not os.path.isdir('./trained_models'):
os.makedirs('./trained_models')
if not os.path.isdir('./trained_models/image_only'):
os.makedirs('./trained_models/image_only')
if not os.path.isdir('./results'):
os.makedirs('./results')
if not os.path.isdir('./results/image_only'):
os.makedirs('./results/image_only')
train_loader = torch.utils.data.DataLoader(datasets.MultiMNIST(
'./data', train=True, download=True, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.MultiMNIST(
'./data', train=False, download=True, transform=transforms.ToTensor()),
batch_size=args.batch_size,
shuffle=True)
vae = ImageVAE(n_latents=args.n_latents)
if args.cuda:
vae.cuda()
optimizer = optim.Adam(vae.parameters(), lr=args.lr)
def train(epoch):
print('Using KL Lambda: {}'.format(kl_lambda))