torch.manual_seed(opts.manual_seed)
if torch.cuda.is_available() and not opts.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
print('Loading data')
pretrain_on_classes = range(5)
if opts.dataset == 'LSUN':
pretrain_on_classes = range(15)
elif opts.dataset == 'Synthetic':
pretrain_on_classes = range(250)
trainset, testset = sup_functions.load_dataset(opts)
indices_train = sup_functions.get_indices_for_classes(trainset,
pretrain_on_classes)
indices_test = sup_functions.get_indices_for_classes(testset,
pretrain_on_classes)
train_loader_classif = data_utils.DataLoader(
trainset,
batch_size=opts.batch_size,
sampler=SubsetRandomSampler(indices_train))
train_loader_gen = data_utils.DataLoader(
trainset,
batch_size=opts.batch_size,
sampler=SubsetRandomSampler(indices_train))
test_loader = data_utils.DataLoader(testset,
batch_size=opts.batch_size,
shuffle=False,
if opts.generator_type == 'AE':
AE_specific = '_' + str(opts.code_size) + '_cl_loss_' + str(opts.betta1) + '_rec_loss_' +str(opts.betta2) +'_'
name_to_save = opts.dataset + '_' + opts.generator_type + AE_specific + str(opts.nb_of_classes) + '_classes.pth'
print(opts)
if opts.manual_seed is None:
opts.manual_seed = random.randint(1, 10000)
print("Random Seed: ", opts.manual_seed)
random.seed(opts.manual_seed)
torch.manual_seed(opts.manual_seed)
if torch.cuda.is_available() and not opts.cuda:
print("WARNING: You have a CUDA device, so you should probably run with --cuda")
print('Loading data')
trainset, testset = sup_functions.load_dataset(opts)
train_loader = data_utils.DataLoader(trainset, batch_size=opts.batch_size, shuffle = True)
test_loader = data_utils.DataLoader(testset, batch_size=opts.batch_size, shuffle = False)
opts.load_classifier = True
classifier = sup_functions.init_classifier(opts)
#classifier.eval()
gen_model = sup_functions.init_generative_model(opts)
criterion_AE = nn.MSELoss()
criterion_classif = nn.MSELoss()
optimizer_gen = torch.optim.Adam(gen_model.parameters(), lr=opts.lr*opts.betta2, betas=(0.9, 0.999), weight_decay=1e-5)
optimizer_classif = torch.optim.Adam(gen_model.parameters(), lr=opts.lr*opts.betta1, betas=(0.9, 0.999), weight_decay=1e-5)
if opts.cuda:
gen_model = gen_model.cuda()
criterion_AE = criterion_AE.cuda()