def main():
opt = get_opt()
print(opt)
print("Start to train stage: %s" % (opt.stage))
# create dataset
if opt.stage == "Shape":
dataset = PolyDatasetShape(128)
train_loader = DataLoader(dataset,
batch_size=opt.b,
shuffle=False,
num_workers=opt.j,
drop_last=True,
pin_memory=True)
elif opt.stage == "Stitch":
dataset = PolyDatasetStitch(128)
train_loader = DataLoader(dataset,
batch_size=opt.b,
shuffle=False,
num_workers=opt.j,
drop_last=True,
pin_memory=True)
elif opt.stage == "Refine":
dataset = PolyDatasetRefine(128)
train_loader = DataLoader(dataset,
batch_size=opt.b,
shuffle=False,
num_workers=opt.j,
drop_last=True,
pin_memory=True)
else:
sys.exit("Please mention the Stage from [Shape, Stitch, Refine]")
if not os.path.exists(opt.results):
os.makedirs(opt.results)
netG = GeneratorCoarse(opt.input_channel, 3)
netD = Discriminator()
# create model & train & save the final checkpoint
netG.cuda()
netD.cuda()
netG.apply(weights_init_normal)
netD.apply(weights_init_normal)
train(opt, train_loader, netG, netD)
print('Finished training %s!' % (opt.stage))
if torch.cuda.is_available():
torch.cuda.manual_seed(opt.seed)
# Networks
if opt.upsample == 'ori':
netG_A2B = Generator_ori(opt.input_nc, opt.output_nc)
netG_B2A = Generator_ori(opt.output_nc, opt.input_nc)
else:
netG_A2B = Generator(opt.input_nc, opt.output_nc)
netG_B2A = Generator(opt.output_nc, opt.input_nc)
netD_A = Discriminator(opt.input_nc)
netD_B = Discriminator(opt.output_nc)
netG_A2B.cuda()
netG_B2A.cuda()
netD_A.cuda()
netD_B.cuda()
netG_A2B.apply(weights_init_normal)
netG_B2A.apply(weights_init_normal)
netD_A.apply(weights_init_normal)
netD_B.apply(weights_init_normal)
torch.save(netG_A2B.state_dict(),
"initial_weights/netG_A2B_seed_{}.pth.tar".format(opt.seed))
torch.save(netG_B2A.state_dict(),
"initial_weights/netG_B2A_seed_{}.pth.tar".format(opt.seed))
torch.save(netD_A.state_dict(),
"initial_weights/netD_A_seed_{}.pth.tar".format(opt.seed))
torch.save(netD_B.state_dict(),
"initial_weights/netD_B_seed_{}.pth.tar".format(opt.seed))
lambda_gp = 10
multi_gpu = False
exp_folder = "{}_{}".format(opt.exp_folder, opt.target_set)
os.makedirs("./exps/" + exp_folder, exist_ok=True)
# Loss function
adversarial_loss = torch.nn.BCEWithLogitsLoss()
distance_loss = torch.nn.L1Loss()
# Initialize generator and discriminator
generator = Generator()
discriminator = Discriminator()
if cuda:
generator.cuda()
discriminator.cuda()
adversarial_loss.cuda()
# Visualize a single batch
def visualizeSingleBatch(fp_loader_test,
opt,
exp_folder,
batches_done,
batch_size=8):
print('Loading saved model ... \n{}'.format(
'./checkpoints/{}_{}.pth'.format(exp_folder, batches_done)))
generatorTest = Generator()
generatorTest.load_state_dict(
torch.load('./checkpoints/{}_{}.pth'.format(exp_folder, batches_done)))
generatorTest = generatorTest.eval()
