random.shuffle(total_random_idx)
imgs = []
labels = []
for idx in total_random_idx:
img, label = data.dataset.__getitem__(idx)
imgs.append(img)
labels.append(torch.tensor(label))
imgs = torch.stack(imgs, dim=0)
labels = torch.stack(labels, dim=0)
return imgs, labels
if args.ckpt_epochs > 0:
if CUDA:
encA.load_state_dict(torch.load('%s/%s-encA_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs)))
decA.load_state_dict(torch.load('%s/%s-decA_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs)))
encB.load_state_dict(torch.load('%s/%s-encB_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs)))
decB.load_state_dict(torch.load('%s/%s-decB_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs)))
else:
encA.load_state_dict(torch.load('%s/%s-encA_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs),
map_location=torch.device('cpu')))
decA.load_state_dict(torch.load('%s/%s-decA_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs),
map_location=torch.device('cpu')))
encB.load_state_dict(torch.load('%s/%s-encB_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs),
map_location=torch.device('cpu')))
decB.load_state_dict(torch.load('%s/%s-decB_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, args.ckpt_epochs),
map_location=torch.device('cpu')))
mask = {}
fixed_imgs = None
for i in range(10):
std.append(torch.std(torch.cat(pred_val[i], dim=0), dim=0))
return torch.stack(std)
def save_ckpt(e):
if not os.path.isdir(args.ckpt_path):
os.mkdir(args.ckpt_path)
torch.save(encA.state_dict(),
'%s/%s-encA_epoch%s.rar' % (args.ckpt_path, MODEL_NAME, e))
pretrain_model = '../weights/mnist_mvae_pretrain/mnist_mvae_pretrain-run_id1-shared10-bs100-lr0.001-encA_epoch140.rar'
if CUDA:
encA.load_state_dict(torch.load(pretrain_model))
else:
encA.load_state_dict(
torch.load(pretrain_model, map_location=torch.device('cpu')))
mask = {}
for e in range(args.ckpt_epochs, args.epochs):
gt_std = get_std(train_data, encA)
train_start = time.time()
train(train_data, encA, e, optimizer, gt_std)
train_end = time.time()
test_start = time.time()
test_accuracy, test_loss = test(test_data, encA)
test_end = time.time()
