def DataIOStream(raw_src: Flow, num: int):
return (raw_src.Take(num).Filter(
lambda x: x.endswith('.jpg')) # select jpg files/选取jpg格式文件
.Map(lambda x: [os.path.join(train_dir, x)] + [
os.path.join(test_dir, x[:-4] + "_" + str(i) + '.jpg')
for i in range(1, 3)
]) # 将噪声数据和真实数据进行合并
.Map(lambda img_file_names: list(
map(
and_then(
data.imread, # 读取图像
img_as_float), # 浮点数张量 [0, 255]->[0, 1]
img_file_names))).Map(to_batch))
def DataIOStream(raw_src: Flow):
return (raw_src.Filter(
lambda x: x.endswith('.jpg')) # select jpg files/选取jpg格式文件
.Map(lambda x: [os.path.join(train_dir, x)] + [
os.path.join(test_dir, x[:-4] + "_" + str(i) + '.jpg')
for i in range(1, 3)
]) # 将噪声数据和真实数据进行合并
.Map(lambda img_file_names: list(
map(
and_then(
data.imread, # 读取图像
img_as_float), # 浮点数张量 [0, 255]->[0, 1]
img_file_names))).Map(to_batch))
train_batches = DataIOStream(raw_sources.Take(train_data_size).Then(cycle))
print('data_loaded')
loss = None
opt = torch.optim.Adam(model.parameters(), lr=lr)
try:
for epoch in range(epochs):
Loss: list = []
train_loss = None
for inner_idx, train in enumerate(
train_batches.Take(train_data_size).Unboxed()):
opt.zero_grad()
train_samples, train_targets = train