def train():
# 加载训练集和验证集
train_loader = load_dataset(opt_train.train_dir,
opt_train.train_size,
opt_train,
shuffled=True)
valid_loader = load_dataset(opt_train.valid_dir,
opt_train.valid_size,
opt_train,
shuffled=False)
# 初始化模型并训练
n2n = Noise2Noise(opt_train, trainable=True)
n2n.train(train_loader, valid_loader)
def test():
# 初始化模型，进行测试
n2n = Noise2Noise(opt_test, trainable=False)
opt_test.redux = False
test_loader = load_dataset(opt_test.data,
3,
opt_test,
shuffled=False,
single=True) #修改0
n2n.load_model(opt_test.load_ckpt) #加载预训练模型
n2n.test(test_loader)
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
if args.pretrained:
model = tf.keras.models.load_model(args.pretrained, compile=False)
else:
# Choose encoder model
model = model_factory.construct_model_from_args(args)
if args.simclr:
model = model_utils.remove_layers(model, CONTRASTIVE_OUTPUT)
# Conrec Dataset (unsupervised)
train, test, num_examples = load_dataset(
args.dataset,
args.data_path,
test_split=args.test_split if not args.no_test_data else None,
train_split=args.train_split,
cache=args.cache,
threads=args.threads,
docker_down=args.docker_download)
conrec_train, conrec_test = (transform.conrec_dataset(
ds,
args.batch_size,
height=args.height,
width=args.width,
implementation=args.aug_impl,
channels=args.channels,
color_jitter_strength=args.color_jitter_strength,
do_shuffle=s,
buffer_multiplier=args.shuffle_buffer_multiplier,
use_blur=args.use_blur) if ds is not None else None
def input_fn():
return load_dataset(config, config.train.batch_size, epochs=-1, shuffle=config.train.shuffle)
def input_fn():
d = load_dataset(config, 1, epochs=1, shuffle=args.shuffle)
d = d.take(args.samples)
return d