train_cfg = cfg["train_data_loader"]
train_zarr = ChunkedDataset(dm.require(train_cfg["key"])).open(
cached=False) # to prevent run out of memory
train_dataset = AgentDataset(cfg, train_zarr, rasterizer)
train_dataloader = DataLoader(train_dataset,
shuffle=train_cfg["shuffle"],
batch_size=train_cfg["batch_size"],
num_workers=train_cfg["num_workers"])
print(train_dataset, len(train_dataset))
train_writer = SummaryWriter('../log/CGAN', comment='CGAN')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 建立模型
generator = cgan.generator(cfg, traj_dim=256, cont_dim=256, mode=3)
discriminator = cgan.discriminator(cfg, h_dim=256, cont_dim=256)
weight_path_g = cfg["model_params"]["weight_path_g"]
if weight_path_g:
generator.load_state_dict(torch.load(weight_path_g))
print(weight_path_g, 'loaded')
weight_path_d = cfg["model_params"]["weight_path_d"]
if weight_path_d:
discriminator.load_state_dict(torch.load(weight_path_d))
print(weight_path_d, 'loaded')
generator.cuda()
discriminator.cuda()
g_loss_fn = utils.g_loss
d_loss_fn = utils.d_loss
learning_rate_g = cfg["model_params"]["lr_g"]
learning_rate_d = cfg["model_params"]["lr_d"]
def train(args):
print(args)
ds_train = VOC2012Dataset(VOC2012_PATH, set='train')
loader_train = data_utils.DataLoader(ds_train,
batch_size=args.batch_size,
num_workers=1,
shuffle=True)
G = Generator(64)
D = discriminator(64)
G.weight_init(mean=0.0, std=0.02)
D.weight_init(mean=0.0, std=0.02)
G.cuda()
D.cuda()
G.train()
D.train()
BCE_Loss = nn.BCELoss().cuda()
L1_Loss = nn.L1Loss().cuda()
print("init_params done.")
if not os.path.exists("./models"):
os.mkdir("./models")
G_optimizer = optim.Adam(G.parameters(),
lr=opt.lrG,
betas=(opt.beta1, opt.beta2))
D_optimizer = optim.Adam(D.parameters(),
lr=opt.lrD,
betas=(opt.beta1, opt.beta2))
real_a = torch.FloatTensor(1, 3, 256, 256)
real_b = torch.FloatTensor(1, 3, 256, 256)
real_a = real_a.cuda()
real_b = real_b.cuda()
real_a = Variable(real_a)
real_b = Variable(real_b)
for epoch in range(opt.train_epoch):
D_losses = []
G_losses = []
epoch_start_time = time.time()
print('training epoch {}'.format(epoch + 1))
for i, (realA, realB) in enumerate(loader_train):
real_a_cpu, real_b_cpu = realA, realB
real_a.data.resize_(real_a_cpu.size()).copy_(real_a_cpu)
real_b.data.resize_(real_b_cpu.size()).copy_(real_b_cpu)
D.zero_grad()
D_result = D(real_a, real_b).squeeze()
D_real_loss = BCE_Loss(
D_result,
torch.Tensor(torch.ones(D_result.size())).cuda())
G_result = G(real_a)
D_result = D(real_a, G_result).squeeze()
D_fake_loss = BCE_Loss(
D_result,
torch.Tensor(torch.zeros(D_result.size())).cuda())
D_train_loss = (D_real_loss + D_fake_loss) * 0.5
D_train_loss.backward()
D_optimizer.step()
G.zero_grad()
G_result = G(real_a)
D_result = D(real_a, G_result).squeeze()
plt.imsave(
'result.png',
(G_result[0].cpu().data.numpy().transpose(1, 2, 0) + 1) / 2)
G_train_loss = BCE_Loss(
D_result,
torch.Tensor(torch.ones(D_result.size())).cuda()
) + opt.L1_lambda * L1_Loss(G_result, real_b)
G_train_loss.backward()
G_optimizer.step()
# if (epoch % 50 == 0):
# torch.save(model, "./models/fcn8s{}.pkl".format(epoch))
# torch.save(netG, "./models/G{}.pkl".format(epoch))
# torch.save(netD, './models/D{}.pkl'.format(epoch))
torch.save(G, "./models/G.pkl")
torch.save(D, './models/D.pkl')
train_hist['total_ptime'].append(total_ptime)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
photo_url = '/data2/limill01/dl/sketchy/photo'
sketch_url = '/data2/limill01/dl/sketchy/sketch'
sketch_photo_pair = ProcessData(photo_url, sketch_url)
split = int(len(sketch_photo_pair) * 0.8)
sketch_photo_pair_train = {}
sketch_photo_pair_test = {}
curr = 0
for key in sketch_photo_pair:
if curr < split:
sketch_photo_pair_train[key] = sketch_photo_pair[key]
else:
sketch_photo_pair_test[key] = sketch_photo_pair[key]
curr += 1
# Params
batch_size = 16
# Instantiate themodel
G = cgan.generator()
D = cgan.discriminator()
G.weight_init(mean=0.0, std=0.02)
D.weight_init(mean=0.0, std=0.02)
training(sketch_photo_pair_train, G, D, batch_size)