def main():
# Step 0: preparation
#place = paddle.fluid.CUDAPlace(0)
with fluid.dygraph.guard():
# Step 1: Define training dataloader
image_folder="work/dummy_data"
image_list_file="work/dummy_data/list.txt"
transform = TrainAugmentation(224)
data = BasicDataLoader(image_folder,image_list_file,transform=transform)
#TODO: create dataloader
train_dataloader = fluid.io.DataLoader.from_generator(capacity=2,return_list=True)
train_dataloader.set_sample_generator(data,args.batch_size)
total_batch = len(data)//args.batch_size
# Step 2: Create model
if args.net == "basic":
#TODO: create basicmodel
model = PSPNet()
else:
raise NotImplementedError(f"args.net: {args.net} is not Supported!")
# Step 3: Define criterion and optimizer
criterion = Basic_SegLoss
# create optimizer
optimizer = AdamOptimizer(learning_rate=args.lr,parameter_list=model.parameters())
# Step 4: Training
for epoch in range(1, args.num_epochs+1):
train_loss = train(train_dataloader,
model,
criterion,
optimizer,
epoch,
total_batch)
print(f"----- Epoch[{epoch}/{args.num_epochs}] Train Loss: {train_loss:.4f}")
if epoch % args.save_freq == 0 or epoch == args.num_epochs:
model_path = os.path.join(args.checkpoint_folder, f"{args.net}-Epoch-{epoch}")
# TODO: save model and optmizer states
model_dict = model.state_dict()
fluid.save_dygraph(model_dict,model_path)
optim_dict = optimizer.state_dict()
fluid.save_dygraph(optim_dict,model_path)
print(f'----- Save model: {model_path}.pdparams')
print(f'----- Save optimizer: {model_path}.pdopt')
def main():
# Step 0: preparation
place = paddle.fluid.CUDAPlace(0)
with fluid.dygraph.guard(place):
# Step 1: Define training dataloader
basic_augmentation = TrainAugmentation(image_size=256)
basic_dataloader = BasicDataLoader(image_folder=args.image_folder,
image_list_file=args.image_list_file,
transform=basic_augmentation,
shuffle=True)
train_dataloader = fluid.io.DataLoader.from_generator(capacity=10, use_multiprocess=True)
train_dataloader.set_sample_generator(basic_dataloader, batch_size=args.batch_size, places=place)
total_batch = int(len(basic_dataloader) / args.batch_size)
# Step 2: Create model
if args.net == "basic":
model = BasicModel()
else:
raise NotImplementedError(f"args.net: {args.net} is not Supported!")
# Step 3: Define criterion and optimizer
criterion = Basic_SegLoss
optimizer = AdamOptimizer(learning_rate=args.lr, parameter_list=model.parameters())
# create optimizer
# Step 4: Training
for epoch in range(1, args.num_epochs+1):
train_loss = train(train_dataloader,
model,
criterion,
optimizer,
epoch,
total_batch)
print(f"----- Epoch[{epoch}/{args.num_epochs}] Train Loss: {train_loss:.4f}")
if epoch % args.save_freq == 0 or epoch == args.num_epochs:
model_path = os.path.join(args.checkpoint_folder, f"{args.net}-Epoch-{epoch}-Loss-{train_loss}")
# TODO: save model and optmizer states
model_dict = model.state_dict()
fluid.save_dygraph(model_dict, model_path)
optimizer_dict = optimizer.state_dict()
fluid.save_dygraph(optimizer_dict, model_path)
print(f'----- Save model: {model_path}.pdparams')
print(f'----- Save optimizer: {model_path}.pdopt')
def main():
# Step 0: preparation
writer = LogWriter(logdir="./log/scalar")
place = paddle.fluid.CUDAPlace(0)
with fluid.dygraph.guard(place):
# Step 1: Define training dataloader
image_folder = ""
image_list_file = "dummy_data/fabric_list.txt"
transform = Transform() #Normalize2() # [0,255]-->[0,1]
x_data = DataLoader(image_folder, image_list_file, transform=transform)
x_dataloader = fluid.io.DataLoader.from_generator(capacity=2,
return_list=True)
x_dataloader.set_sample_generator(x_data, args.batch_size)
total_batch = len(x_data) // args.batch_size
# Step 2: Create model
if args.net == "basic":
D = Discriminator()
G = Generator()
E = Invertor()
else:
raise NotImplementedError(
f"args.net: {args.net} is not Supported!")
# Step 3: Define criterion and optimizer
criterion = Basic_Loss
D_optim = AdamOptimizer(learning_rate=args.lr,
parameter_list=D.parameters())
G_optim = AdamOptimizer(learning_rate=args.lr,
parameter_list=G.parameters())
E_optim = AdamOptimizer(learning_rate=args.lr,
parameter_list=E.parameters())
G_loss_meter = AverageMeter()
D_loss_meter = AverageMeter()
E_loss_meter = AverageMeter()
D.train()
G.train()
E.train()
# Step 4: Slight Training
iteration = -1
is_slight_Train = True
for epoch in range(1, args.epoch_num + 1):
#optim Discriminator
for (x, x_labels) in x_dataloader():
n = x.shape[0]
if is_slight_Train:
iteration += 1
x = fluid.layers.cast(x, dtype="float32")
x = fluid.layers.transpose(x, perm=[0, 3, 1, 2])
preds_x = D(x)
preds_x_array = preds_x.numpy()
#print("D(x),1",preds_array.shape, np.mean(preds_array))
writer.add_scalar(tag="D(x)=1",
step=iteration,
value=np.mean(preds_x_array))
if np.mean(preds_x_array) >= 0.98:
is_slight_Train = False
z = np.random.rand(n, 64)
zeros = np.zeros((n, 1))
z = to_variable(z)
zeros = to_variable(zeros)
z = fluid.layers.cast(z, dtype="float32")
zeros = fluid.layers.cast(zeros, dtype="int64")
preds_fx = D(G(z))
preds_fx_array = preds_fx.numpy()
writer.add_scalar(tag="D(G(z))=0",
step=iteration,
value=np.mean(preds_fx_array))
D_loss = criterion(preds_x, x_labels) + criterion(
preds_fx, zeros)
D_loss.backward()
D_optim.minimize(D_loss)
D.clear_gradients()
D_loss_meter.update(D_loss.numpy()[0], n)
writer.add_scalar(tag="D_loss",
step=iteration,
value=D_loss_meter.avg)
print(f"EPOCH[{epoch:03d}/{args.epoch_num:03d}], " +
f"STEP{iteration}, " +
f"Average D Loss: {D_loss_meter.avg:4f}, ")
z = np.random.rand(n, 64)
ones = np.ones((n, 1))
z = to_variable(z)
ones = to_variable(ones)
z = fluid.layers.cast(z, dtype="float32")
ones = fluid.layers.cast(ones, dtype="int64")
preds = D(G(z))
preds_array = preds.numpy()
writer.add_scalar(tag="D(G(z))=1",
step=iteration,
value=np.mean(preds_array))
G_loss = criterion(preds, ones)
G_loss.backward()
G_optim.minimize(G_loss)
G.clear_gradients()
G_loss_meter.update(G_loss.numpy()[0], n)
writer.add_scalar(tag="G_loss",
step=iteration,
value=G_loss_meter.avg)
print(f"EPOCH[{epoch:03d}/{args.epoch_num:03d}], " +
f"STEP{iteration}, " +
f"Average G Loss: {G_loss_meter.avg:4f}")
if epoch % args.save_freq == 0 or epoch == args.epoch_num or not is_slight_Train:
D_model_path = os.path.join(args.checkpoint_folder,
f"D_{args.net}-Epoch-{epoch}")
G_model_path = os.path.join(args.checkpoint_folder,
f"G_{args.net}-Epoch-{epoch}")
# save model and optmizer states
model_dict = D.state_dict()
fluid.save_dygraph(model_dict, D_model_path)
optim_dict = D_optim.state_dict()
fluid.save_dygraph(optim_dict, D_model_path)
model_dict = G.state_dict()
fluid.save_dygraph(model_dict, G_model_path)
optim_dict = G_optim.state_dict()
fluid.save_dygraph(optim_dict, G_model_path)
print(
f'----- Save model: {D_model_path}.pdparams, {G_model_path}.pdparams'
)
if not is_slight_Train:
break
# Step 5: full training for Generator and Discriminator
D_optim = AdamOptimizer(learning_rate=args.lr * 10,
parameter_list=D.parameters())
G_optim = AdamOptimizer(learning_rate=args.lr * 10,
parameter_list=G.parameters())
G_loss_meter = AverageMeter()
D_loss_meter = AverageMeter()
for epoch in range(1, args.epoch_num + 1):
for (x, x_labels) in x_dataloader():
n = x.shape[0]
iteration += 1
x = fluid.layers.cast(x, dtype="float32")
x = fluid.layers.transpose(x, perm=[0, 3, 1, 2])
preds1 = D(x)
preds_array = preds1.numpy()
writer.add_scalar(tag="D(x)=1",
step=iteration,
value=np.mean(preds_array))
z = np.random.rand(n, 64)
zeros = np.zeros((n, 1))
z = to_variable(z)
zeros = to_variable(zeros)
z = fluid.layers.cast(z, dtype="float32")
zeros = fluid.layers.cast(zeros, dtype="int64")
preds2 = D(G(z))
preds_array = preds2.numpy()
#print("DG(z),0:",preds_array.shape, np.mean(preds_array))
writer.add_scalar(tag="D(G(z))=0",
step=iteration,
value=np.mean(preds_array))
D_loss = criterion(preds1, x_labels) + criterion(preds2, zeros)
D_loss.backward()
D_optim.minimize(D_loss)
D.clear_gradients()
D_loss_meter.update(D_loss.numpy()[0], n)
writer.add_scalar(tag="D_loss",
step=iteration,
value=D_loss_meter.avg)
print(f"EPOCH[{epoch:03d}/{args.epoch_num:03d}], " +
f"STEP{iteration}, " +
f"Average D Loss: {D_loss_meter.avg:4f} ")
z = np.random.rand(n, 64)
ones = np.ones((n, 1))
z = to_variable(z)
ones = to_variable(ones)
z = fluid.layers.cast(z, dtype="float32")
ones = fluid.layers.cast(ones, dtype="int64")
preds = D(G(z))
preds_array = preds.numpy()
#print("DG(z),1:",preds_array.shape, np.mean(preds_array))
writer.add_scalar(tag="D(G(z))=1",
step=iteration,
value=np.mean(preds_array))
G_loss = criterion(preds, ones)
G_loss.backward()
G_optim.minimize(G_loss)
G.clear_gradients()
G_loss_meter.update(G_loss.numpy()[0], n)
writer.add_scalar(tag="G_loss",
step=iteration,
value=G_loss_meter.avg)
print(f"EPOCH[{epoch:03d}/{args.epoch_num:03d}], " +
f"STEP{iteration}, " +
f"Average G Loss: {G_loss_meter.avg:4f}")
if epoch % args.save_freq == 0 or epoch == args.epoch_num:
D_model_path = os.path.join(args.checkpoint_folder,
f"D_{args.net}-Epoch-{epoch}")
G_model_path = os.path.join(args.checkpoint_folder,
f"G_{args.net}-Epoch-{epoch}")
# save model and optmizer states
model_dict = D.state_dict()
fluid.save_dygraph(model_dict, D_model_path)
optim_dict = D_optim.state_dict()
fluid.save_dygraph(optim_dict, D_model_path)
model_dict = G.state_dict()
fluid.save_dygraph(model_dict, G_model_path)
optim_dict = G_optim.state_dict()
fluid.save_dygraph(optim_dict, G_model_path)
print(
f'----- Save model: {D_model_path}.pdparams, {G_model_path}.pdparams'
)
# Step 6: full training for Inverter
E_optim = AdamOptimizer(learning_rate=args.lr * 10,
parameter_list=E.parameters())
E_loss_meter = AverageMeter()
for epoch in range(1, args.epoch_num + 1):
for (x, x_labels) in x_dataloader():
n = x.shape[0]
iteration += 1
x = fluid.layers.cast(x, dtype="float32")
image = x.numpy()[0] * 255
writer.add_image(tag="x", step=iteration, img=image)
x = fluid.layers.transpose(x, perm=[0, 3, 1, 2])
invert_x = G(E(x))
invert_image = fluid.layers.transpose(invert_x,
perm=[0, 2, 3, 1])
invert_image = invert_image.numpy()[0] * 255
#print("D(x),1",preds_array.shape, np.mean(preds_array))
writer.add_image(tag="invert_x",
step=iteration,
img=invert_image)
print(np.max(invert_image), np.min(invert_image))
E_loss = fluid.layers.mse_loss(invert_x, x)
print("E_loss shape:", E_loss.numpy().shape)
E_loss.backward()
E_optim.minimize(E_loss)
E.clear_gradients()
E_loss_meter.update(E_loss.numpy()[0], n)
writer.add_scalar(tag="E_loss",
step=iteration,
value=E_loss_meter.avg)
print(f"EPOCH[{epoch:03d}/{args.epoch_num:03d}], " +
f"STEP{iteration}, " +
f"Average E Loss: {E_loss_meter.avg:4f}, ")
if epoch % args.save_freq == 0 or epoch == args.epoch_num:
E_model_path = os.path.join(args.checkpoint_folder,
f"E_{args.net}-Epoch-{epoch}")
# save model and optmizer states
model_dict = E.state_dict()
fluid.save_dygraph(model_dict, E_model_path)
optim_dict = E_optim.state_dict()
fluid.save_dygraph(optim_dict, E_model_path)
print(
f'----- Save model: {E_model_path}.pdparams, {E_model_path}.pdparams'
)