# 创建验证模型
model_val = create_model(opt_val)
model_val.eval()
# 训练设置
opt_train = TrainOptions().parse()
# 设置显示训练结果的类
visualizer = Visualizer(opt_train)
for epoch in range(opt_train.epoch_count,
opt_train.niter + opt_train.niter_decay + 1):
epoch_iters = 0
epoch_start_time = time.time()
# 验证结果
metrics = RunningScore(opt_val.num_classes)
model_val.opt.epoch = epoch
model_val.setup(model_val.opt)
for i, data in enumerate(dataset_val):
model_val.set_input(data)
model_val.forward()
gt = np.squeeze(data["label"].numpy(), axis=1) # [N, W, H]
pre = model_val.pre
pre = pre.data.max(1)[1].cpu().numpy() # [N, W, H]
metrics.update(gt, pre)
# 保存结果
if i % opt_train.display_freq == 0: # 逻辑有点问题
save_segment_result(model_val.get_current_visuals(), epoch,
opt_train.display_winsize, image_dir,
# 创建训练模型
model_train = create_model(opt_train)
model_train.train()
opt_train.continue_train = True
model_train.setup(opt_train)
# 创建验证模型
opt_val.isTrain = False
model_val = create_model(opt_val)
model_val.eval()
# 设置显示训练结果的类
visualizer = Visualizer(opt_train)
metrics = RunningScore(opt_train.num_classes)
for epoch in range(opt_train.epoch_count,
opt_train.niter + opt_train.niter_decay + 1):
epoch_iters = 0
epoch_start_time = time.time()
metrics.reset()
for i, data in enumerate(dataset_train):
# label, image.shape == [N, C, W, H]
iter_start_time = time.time()
epoch_iters += 1
# 训练一次
model_train.set_input(data)
model_train.optimize_parameters()
opt_val = TestOptions().parse()
# 加载验证数据集
opt_val.dataset_mode = 'single'
opt_val.batch_size = 1
dataset_val = create_dataset(opt_val)
dataset_val_size = len(dataset_val)
print('The number of valling images = %d' % dataset_val_size)
# 创建验证模型
model_val = create_model(opt_val)
model_val.eval()
# 从这一轮保存的权重恢复
model_val.setup(opt_val)
metrics = RunningScore(opt_val.num_classes)
metrics.reset()
for i, data in enumerate(dataset_val):
model_val.set_input(data)
model_val.forward()
gt = data["B_label"].numpy().squeeze() # [H, W]
output = model_val.pre # [N, C, H, W]
output = nn.functional.softmax(output, dim=1) # [N, C, H, W]
output = nn.functional.upsample(
output, (1024, 2048), mode='bilinear',
align_corners=True).cpu().data[0].numpy() # [C, H, W]
output = output.transpose(1, 2, 0) # [H, W, C]
output_nomask = np.asarray(np.argmax(output, axis=2),
dtype=np.uint8) # [H, W]
# 保存成彩色图