def build_model(self):
"""
Instantiates the model, loss criterion, and optimizer
"""
# instantiate model
self.model = VGGNet(self.config, self.use_batch_norm,
self.input_channels, self.class_count,
self.init_weights)
# instantiate loss criterion
self.criterion = nn.CrossEntropyLoss()
# instantiate optimizer
self.optimizer = optim.SGD(self.model.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay)
# print network
self.print_network(self.model, 'VGGNet')
# use gpu if enabled
if torch.cuda.is_available() and self.use_gpu:
self.model.cuda()
self.criterion.cuda()
def train(**kwargs):
cfg = Config()
for k, v in kwargs.items():
setattr(cfg, k, v)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
content = load_image(cfg.content, transform, max_size=cfg.max_size)
style = load_image(cfg.style,
transform,
shape=[content.size(3),
content.size(2)])
target = Variable(content.clone(), requires_grad=True)
optimizer = torch.optim.Adam([target], lr=cfg.lr, betas=[0.5, 0.999])
vgg = VGGNet()
if cfg.use_gpu:
vgg.cuda()
for step in range(cfg.total_step):
target_features = vgg(target)
content_features = vgg(Variable(content))
style_features = vgg(Variable(style))
style_loss = 0
content_loss = 0
for f1, f2, f3 in zip(target_features, content_features,
style_features):
# Compute content loss
content_loss += torch.mean((f1 - f2)**2)
_, c, h, w = f1.size()
f1 = f1.view(c, h * w)
f3 = f3.view(c, h * w)
# Compute gram matrix
f1 = torch.mm(f1, f1.t())
f3 = torch.mm(f3, f3.t())
style_loss += torch.mean((f1 - f3)**2) / (c * h * w)
loss = content_loss + cfg.style_weight * style_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step + 1) % cfg.log_step == 0:
print('Step [%d/%d], Content Loss: %.4f, Style Loss: %.4f' %
(step + 1, cfg.total_step, content_loss.data[0],
style_loss.data[0]))
if (step + 1) % cfg.sample_step == 0:
denorm = transforms.Normalize((-2.12, -2.04, -1.80),
(4.37, 4.46, 4.44))
img = target.clone().cpu().squeeze()
img = denorm(img.data).clamp_(0, 1)
torchvision.utils.save_image(img, 'output-%d.png' % (step + 1))
def train():
db = './public/cards/trainImg'
img_list = get_imlist(db, '.jpg')
model = VGGNet()
feats = []
names = []
for i, img_path in enumerate(img_list):
norm_feat = model.extract_feat(img_path)
img_name = os.path.split(img_path)[1]
feats.append(norm_feat)
names.append(img_name)
print("extracting feature from image No. %d , %d images in total" %
((i+1), len(img_list)))
feats = np.array(feats)
output = './src/vis/model/model.h5'
h5f = h5py.File(output, 'w')
h5f.create_dataset('dataset_1', data=feats)
h5f.create_dataset('dataset_2', data=np.string_(names))
h5f.close()
def predict(img_path, weights_path, class_indices_path):
"""
对图像进行预测分类
:param img_path: 待预测图像路径
:param weights_path: 模型权重路径
:param class_indices_path: 标签类别索引
:return: 待预测图像类别
"""
img_height = img_width = 224
# 加载待预测图像
img = Image.open(img_path)
# 重设图像大小
img = img.resize((img_width, img_height))
plt.imshow(img)
# 归一化
img = np.array(img) / 255.
# 增加batch这个维度
img = (np.expand_dims(img, 0))
# 加载标签类别索引文件
try:
json_file = open(class_indices_path, 'r')
class_indict = json.load(json_file)
except Exception as e:
print(e)
exit(-1)
# 预测
model = VGGNet(img_height, img_width, class_num=5, name='vgg11').vgg()
model.load_weights(weights_path)
result = np.squeeze(model.predict(img))
predict_class = np.argmax(result)
label = class_indict[str(predict_class)], result[predict_class]
plt.title(label)
plt.show()
size = np.array(content_image.size) * scale
content_image = content_image.resize(size.astype(int), Image.ANTIALIAS)
content_image = transform(content_image).unsqueeze(0).cuda()
# Style Image processing
style_image = Image.open(style_image)
style_image = style_image.resize(
[content_image.size(2), content_image.size(3)], Image.LANCZOS)
style_image = transform(style_image).unsqueeze(0).cuda()
# Initialize result and optimizer
result_image = Variable(content_image.clone(), requires_grad=True)
optimizer = torch.optim.Adam([result_image], lr=0.003, betas=[0.5, 0.999])
# Model
vgg = VGGNet()
vgg = vgg.cuda()
# Train
for step in range(epochs):
target_features = vgg(result_image)
content_features = vgg(Variable(content_image))
style_features = vgg(Variable(style_image))
style_loss = 0
content_loss = 0
for f1, f2, f3 in zip(target_features, content_features, style_features):
# Content loss
content_loss += torch.mean((f1 - f2)**2)
import tensorflow as tf
import sys
from model import VGGNet
from data_loader import DataLoader
net = VGGNet([224, 224], 128)
net.build()
loss = net.loss()
# print(tf.global_variables())
ckpt_path = '../ckpt/model.ckpt-0'
loader = DataLoader()
sess = tf.Session()
optimizer = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(loss)
saver = tf.train.Saver(tf.global_variables(), max_to_keep=None)
ls = tf.summary.scalar('loss', loss)
train_writer = tf.summary.FileWriter('../log_train', sess.graph)
valid_writer = tf.summary.FileWriter('../log_valid', sess.graph)
batch = 32
batch_num = loader.images_urls.shape[0] // batch
# config = tf.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction = 0.7
valid_batch_num = loader.valid_urls.shape[0] // batch
if ckpt_path:
saver.restore(sess, ckpt_path)
else:
sess.run(tf.global_variables_initializer())
from model import VGGNet
batch_size=32
train_transforms = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
test_dir = './Atestdataset'
test_datasets = datasets.ImageFolder(test_dir, transform=train_transforms)
test_dataloader = torch.utils.data.DataLoader(test_datasets, batch_size=batch_size, shuffle=True)
model = VGGNet()
model.load_state_dict(torch.load('./Vgg16.pth',map_location=torch.device('cpu')))
for epoch in range(1):
model = model.eval()
total = 0
correct = 0
for i, data in enumerate(test_dataloader):
images, labels = data
vggoutputs = model(images)
_, vggpredicted = torch.max(vggoutputs.data, 1)
#print(labels)
#print(vggpredicted)
total += labels.size(0)
correct += (vggpredicted == labels).sum().item()
print(100.0 * correct / total)
def run(model_name):
train_dir = '../flower_data/train'
validation_dir = '../flower_data/validation'
if not os.path.exists('../save_weights'):
os.mkdir('../save_weights')
img_height, img_width = 224, 224
batch_size = 32
epochs = 10
# 准备训练集和验证集
train_image_generator = ImageDataGenerator(rescale=1. / 255, horizontal_flip=True)
validation_image_generator = ImageDataGenerator(rescale=1. / 255)
train_data_gen = train_image_generator.flow_from_directory(directory=train_dir,
batch_size=batch_size,
shuffle=True,
target_size=(img_height, img_width),
class_mode='categorical')
valid_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir,
batch_size=batch_size,
shuffle=False,
target_size=(img_height, img_width),
class_mode='categorical')
# 训练样本数和验证样本数
total_train, total_valid = train_data_gen.n, valid_data_gen.n
# 生成标签索引字典
class_indices = train_data_gen.class_indices
inverse_dict = dict((v, k) for k, v in class_indices.items())
json_str = json.dumps(inverse_dict, indent=4)
with open('../class_indices.json', 'w') as json_file:
json_file.write(json_str)
# 训练模型
model = VGGNet(img_height, img_width, class_num=5, name=model_name).vgg()
model.summary()
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0005),
loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),
metrics=['accuracy'])
callbacks = [tf.keras.callbacks.ModelCheckpoint(filepath=f'../save_weights/{model_name}.h5',
save_best_only=True,
save_weights_only=True,
monitor='val_loss')]
history = model.fit(x=train_data_gen,
steps_per_epoch=total_train // batch_size,
epochs=epochs,
validation_data=valid_data_gen,
validation_steps=total_valid // batch_size,
callbacks=callbacks)
# 评估模型
# plot loss and accuracy image
history_dict = history.history
train_loss = history_dict["loss"]
train_accuracy = history_dict["accuracy"]
val_loss = history_dict["val_loss"]
val_accuracy = history_dict["val_accuracy"]
# figure 1
plt.figure()
plt.plot(range(epochs), train_loss, label='train_loss')
plt.plot(range(epochs), val_loss, label='val_loss')
plt.legend()
plt.xlabel('epochs')
plt.ylabel('loss')
# figure 2
plt.figure()
plt.plot(range(epochs), train_accuracy, label='train_accuracy')
plt.plot(range(epochs), val_accuracy, label='val_accuracy')
plt.legend()
plt.xlabel('epochs')
plt.ylabel('accuracy')
plt.show()
import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
from dataset import CityScapeDataset
from model import FCNs, VGGNet
from torchvision import transforms, utils
from torch import Tensor
from labels import *
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
batch_size, n_class, h, w = 1, 19, 480, 320
vgg_model = VGGNet(requires_grad=True)
fcn_model = FCNs(pretrained_net=vgg_model, n_class=n_class).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(fcn_model.parameters(), lr=1e-3, momentum=0.9)
preprocess = transforms.Compose([
# transforms.Scale(256),
# transforms.ToTensor(),
# normalize
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
dataset = CityScapeDataset('.\\data', (h, w),
transform=preprocess,
target_transform=transforms.RandomHorizontalFlip())
img, label = dataset[0]
import tensorflow as tf
from model import VGGNet
import numpy as np
from data_loader import DataLoader
import time
ckpt_path = '../ckpt/model_0.ckpt'
net = VGGNet([224, 224], 128, training=False)
net.build()
sess = tf.Session()
saver = tf.train.Saver(tf.global_variables(), max_to_keep=None)
# config = tf.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction = 0.7
if ckpt_path:
saver.restore(sess, ckpt_path)
loader = DataLoader()
batch = 64
valid_batch_num = loader.valid_urls.shape[0] // batch
cou = 1
# for idx in range(valid_batch_num):
# res = loader.get_valid_batch_data(batch)
# feed_dicts = {net.inputs: res[0], net.ground_truth: res[1]}
# # sess.run(optimizer, feed_dict=feed_dicts)
# fc_16 = sess.run([net.fc_16], feed_dict=feed_dicts)
# fc_16 = np.array(fc_16[0])
# for i in range(batch):
# if np.argmax(fc_16[i, :]) == np.argmax([res[1][i, :]]):
import torch
import torch.nn as nn
import torch.optim as optim
from model import VGGNet, VGG_CONFS
MODEL_PATH = 'models/checkpoint_e49.pkl'
SAMPLE_IMG_PATH = ''
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
cpt = torch.load(MODEL_PATH, map_location=device)
epoch = cpt['epoch']
seed = cpt['seed']
total_steps = cpt['total_steps']
# make even this the same...
vgg16 = nn.parallel.DataParallel(
VGGNet(VGG_CONFS['vgg16'], dim=32, num_classes=10))
vgg16.load_state_dict(cpt['model'])
print(vgg16)
# test loading optimizer
optimizer = optim.SGD(vgg16.parameters(),
lr=0.0001,
weight_decay=0.00005,
momentum=0.9).load_state_dict(cpt['optimizer'])
download=True,
transform=transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(32,
scale=(0.8,
1.0)),
transforms.ToTensor(),
normalize,
]))
print('Dataset created - size: {}'.format(len(dataset)))
seed = torch.initial_seed()
print('Using seed : {}'.format(seed))
# create model & train on multiple GPUs
vggnet = VGGNet(VGG_CONFS[MODEL_TYPE], dim=IMAGE_DIM,
num_classes=NUM_CLASSES).to(device)
vggnet = torch.nn.parallel.DataParallel(vggnet, device_ids=DEVICE_IDS)
print(vggnet)
print('VGGNet created')
dataloader = data.DataLoader(dataset,
shuffle=True,
pin_memory=True,
drop_last=True,
num_workers=4,
batch_size=BATCH_SIZE)
print('Dataloader created')
# create optimizer
optimizer = optim.SGD(params=vggnet.parameters(),
lr=LR_INIT,
def train(cfg):
n_class = int(cfg["data"]["n_class"])
img_h = int(cfg["data"]["img_h"])
img_w = int(cfg["data"]["img_w"])
batch_size = int(cfg["training"]["batch_size"])
epochs = int(cfg["training"]["epochs"])
lr = float(cfg["training"]["optimizer"]["lr"])
momentum = float(cfg["training"]["optimizer"]["momentum"])
w_decay = float(cfg["training"]["optimizer"]["weight_decay"])
step_size = int(cfg["training"]["lr_schedule"]["step_size"])
gamma = float(cfg["training"]["lr_schedule"]["gamma"])
configs = "FCNs-BCEWithLogits_batch{}_epoch{}_RMSprop_scheduler-step{}-gamma{}_lr{}_momentum{}_w_decay{}_input_size{}_03091842".format(
batch_size, epochs, step_size, gamma, lr, momentum, w_decay, img_h)
print("Configs:", configs)
root_dir = cfg["data"]["root_dir"]
train_filename = cfg["data"]["train_file"]
val_filename = cfg["data"]["val_file"]
mean_filename = cfg["data"]["mean_file"]
class_weight_filename = cfg["data"]["class_weight_file"]
train_file = os.path.join(root_dir, train_filename)
print(train_file)
val_file = os.path.join(root_dir, val_filename)
mean_file = os.path.join(root_dir, mean_filename)
class_weight_file = os.path.join(root_dir, class_weight_filename)
model_dir = cfg["training"]["model_dir"]
if not os.path.exists(model_dir):
os.makedirs(model_dir)
model_path = os.path.join(model_dir, configs)
use_gpu = torch.cuda.is_available()
num_gpu = list(range(torch.cuda.device_count()))
continue_train = False
#MeanRGB_train = ComputeMeanofInput(train_file)
#MeanRGB_train = np.load(mean_file)
MeanRGB_train = np.array([0.0, 0.0, 0.0])
print("MeanRGB_train: {}".format(MeanRGB_train))
train_data = ScanNet2d(csv_file=train_file,
phase='train',
trainsize=(img_h, img_w),
MeanRGB=MeanRGB_train)
val_data = ScanNet2d(csv_file=val_file,
phase='val',
trainsize=(img_h, img_w),
MeanRGB=MeanRGB_train)
train_loader = DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
num_workers=1)
val_loader = DataLoader(val_data,
batch_size=batch_size,
shuffle=False,
num_workers=1)
#class_weight = trainer.computer_class_weights(train_file)
class_weight = np.load(class_weight_file)
print("class_weight: {}".format(class_weight))
class_weight = torch.from_numpy(class_weight)
print("shape of class weight {}".format(class_weight.shape))
vgg_model = VGGNet(requires_grad=True, remove_fc=True)
fcn_model = FCN8s(encoder_net=vgg_model, n_class=n_class)
if use_gpu:
ts = time.time()
vgg_model = vgg_model.cuda()
fcn_model = fcn_model.cuda()
fcn_model = nn.DataParallel(fcn_model, device_ids=num_gpu)
class_weight = class_weight.cuda()
print("Finish cuda loading, tme elapsed {}".format(time.time() - ts))
L = nn.BCEWithLogitsLoss(reduction='none')
optimizer = optim.RMSprop(fcn_model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=w_decay)
#optimizer = optim.SGD(fcn_model.parameters(), lr=lr, momentum= momentum, weight_decay=w_decay)
scheduler = lr_scheduler.StepLR(optimizer,
step_size=step_size,
gamma=gamma)
score_dir = os.path.join("scores", configs)
if not os.path.exists(score_dir):
os.makedirs(score_dir)
log_headers = [
'epoch', 'train/loss', 'train/acc', 'train/acc_cls', 'train/mean_iu',
'train/fwavacc', 'val/loss', 'val/acc', 'val/acc_cls', 'val/mean_iu',
'val/fwavacc', 'elapsed_time'
]
if not os.path.exists(os.path.join(score_dir, 'log.csv')):
with open(os.path.join(score_dir, 'log.csv'), 'w') as f:
f.write(','.join(log_headers) + '\n')
IU_scores = np.zeros((epochs, n_class + 1))
pixel_scores = np.zeros(epochs)
writer = SummaryWriter()
# color_mapping = util.GenerateColorMapping(n_class)
best_mean_iu = 0
epoch_loss = 0.0
if continue_train:
model_path = "C:\\Users\\ji\\Documents\\FCN-VGG16\\models\\FCNs-BCEWithLogits_batch1_epoch500_RMSprop_scheduler-step50-gamma0.5_lr0.0001_momentum0.0_w_decay1e-05"
fcn_model = torch.load(model_path)
fcn_model.train()
for epoch in range(epochs):
fcn_model.train()
scheduler.step()
ts = time.time()
running_loss = 0.0
######
label_preds = []
label_trues = []
######
for i, batch in enumerate(train_loader):
optimizer.zero_grad()
if use_gpu:
inputs = Variable(batch['X'].cuda())
labels = Variable(batch['Y'].cuda())
else:
inputs, labels = Variable(batch['X']), Variable(batch['Y'])
outputs = fcn_model(inputs)
#print("out: {}".format(outputs.shape))
#print("label: {}".format(labels.shape))
#print(outputs)
#print(labels)
loss = L(outputs, labels)
# print(loss.shape)
loss = loss.permute(0, 2, 3,
1).reshape(-1,
n_class + 1) #.view(-1,n_class+1)
# print(loss.shape)
loss = torch.mean(loss, dim=0)
# print(loss.shape)
loss = torch.mul(loss, class_weight)
# print(loss.shape)
loss = torch.mean(loss)
# print(loss)
loss.backward()
# print("grad")
# print(fcn_model.outp.weight.grad)
# print(fcn_model.embs[0].weight.grad)
optimizer.step()
#scheduler.step()
if i == 0 and epoch == 0:
# count= util.count_parameters(fcn_model)
# print("number of parameters in model {}".format(count))
visIn = inputs[:3]
#print('shape of in {}'.format(visIn[:5].shape))
visLabel = batch['l'][:3]
epoch_loss += loss.item()
running_loss += loss.item()
# print("loss: {}".format(loss.data))
if i % 10 == 9 and i != 0:
print("epoch{}, iter{}, Iterloss: {}".format(
epoch, i, running_loss / 10))
writer.add_scalar('train/iter_loss', running_loss / 10,
epoch * len(train_loader) + i)
running_loss = 0.0
# N, _, h, w = outputs.shape
# targets = batch['l'].cpu().numpy().reshape(N,h,w)
# outputs = outputs.data.cpu().numpy()
# preds_v, targets_v = util.visulaize_output(outputs,targets,color_mapping,n_class)
# writer.add_images('train/predictions',torch.from_numpy(preds_v),dataformats='NHWC')
# writer.add_images('train/targets',torch.from_numpy(targets_v),dataformats='NHWC')
#####################################
outputs = outputs.data.cpu().numpy()
N, _, h, w = outputs.shape
pred = outputs.transpose(0, 2, 3, 1).reshape(
-1, n_class + 1).argmax(axis=1).reshape(N, h, w)
target = batch['l'].cpu().numpy().reshape(N, h, w)
#########
for lt, lp in zip(target, pred):
label_trues.append(lt)
label_preds.append(lp)
metrics = util.label_accuracy_score(label_trues, label_preds,
n_class + 1)
with open(os.path.join(score_dir, "log.csv"), 'a') as f:
log = [epoch] + [epoch_loss] + list(metrics) + [''] * 7
log = map(str, log)
f.write(','.join(log) + '\n')
########################################
#scheduler.step()
writer.add_scalar('train/epoch_loss', epoch_loss, epoch)
print("Finish epoch{}, epoch loss {}, time eplapsed {}".format(
epoch, epoch_loss,
time.time() - ts))
epoch_loss = 0.0
####################
writer.add_scalar('train/mean_iu', metrics[2], epoch)
writer.add_scalar('train/acc', metrics[0], epoch)
writer.add_scalar('train/acc_cls', metrics[1], epoch)
######################
#Training precess visulize
visOut = fcn_model(visIn)
preds_v, targets_v = util.visulaize_output(visOut, visLabel, n_class)
writer.add_images('train/predictions',
torch.from_numpy(preds_v),
global_step=epoch,
dataformats='NHWC')
writer.add_images('train/targets',
torch.from_numpy(targets_v),
global_step=epoch,
dataformats='NHWC')
if not os.path.exists(model_path):
os.makedirs(model_path)
torch.save(fcn_model, os.path.join(model_path, str(epoch)))
best_mean_iu = val_model(epoch, val_loader, fcn_model, use_gpu,
n_class, IU_scores, pixel_scores, score_dir,
writer, best_mean_iu, model_path, L)
writer.flush()
writer.close()
def main():
# 生成一个图像,均值为127.5，方差为20
result = initial_result((1, 466, 712, 3), 127.5, 20)
# 读取内容图像和风格图像
content_val = read_img(content_img_path)
style_val = read_img(style_img_path)
content = tf.placeholder(tf.float32, shape=[1, 466, 712, 3])
style = tf.placeholder(tf.float32, shape=[1, 615, 500, 3])
# 载入模型，注意：在python3中，需要添加一句： encoding='latin1'
data_dict = np.load(vgg_16_npy_pyth, encoding='latin1').item()
# 创建这三张图像的 vgg 对象
vgg_for_content = VGGNet(data_dict)
vgg_for_style = VGGNet(data_dict)
vgg_for_result = VGGNet(data_dict)
# 创建每个神经网络
vgg_for_content.build(content)
vgg_for_style.build(style)
vgg_for_result.build(result)
# 提取哪些层特征
# 需要注意的是：内容特征抽取的层数和结果特征抽取的层数必须相同
# 风格特征抽取的层数和结果特征抽取的层数必须相同
content_features = [
# vgg_for_content.conv1_2,
# vgg_for_content.conv2_2,
# vgg_for_content.conv3_3,
vgg_for_content.conv4_3,
vgg_for_content.conv5_3,
]
result_content_features = [
# vgg_for_result.conv1_2,
# vgg_for_result.conv2_2,
# vgg_for_result.conv3_3,
vgg_for_result.conv4_3,
vgg_for_result.conv5_3,
]
style_features = [
vgg_for_style.conv2_2,
]
result_style_features = [
vgg_for_result.conv2_2,
]
style_gram = [gram_matrix(feature) for feature in style_features]
result_style_gram = [
gram_matrix(feature) for feature in result_style_features
]
# 计算内容损失
content_loss = tf.zeros(1, tf.float32)
for c, c_ in zip(content_features, result_content_features):
content_loss += tf.reduce_mean((c - c_)**2, axis=[1, 2, 3])
# 计算风格损失
style_loss = tf.zeros(1, tf.float32)
for s, s_ in zip(style_gram, result_style_gram):
# 因为在计算gram矩阵的时候，降低了一维，所以，只需要在[1, 2]两个维度求均值即可
style_loss += tf.reduce_mean((s - s_)**2, axis=[1, 2])
# 总的损失函数
loss = content_loss * lambda_c + style_loss * lambda_s
train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
print("1111111")
for step in range(num_steps):
loss_value, content_loss_value, style_loss_value, _ = sess.run(
[loss, content_loss, style_loss, train_op],
feed_dict={
content: content_val,
style: style_val
})
print(
'step: %d, loss_value: %.4f, content_loss: %.4f, style_loss: %.4f'
% (step + 1, loss_value[0], content_loss_value[0],
style_loss_value[0]))
if step % 100 == 0:
result_img_path = os.path.join(output_dir,
'result_%05d.jpg' % (step + 1))
# 将图像取出，因为之前是4维，所以需要使用一个索引0，将其取出
result_val = result.eval(sess)[0]
# np.clip() numpy.clip(a, a_min, a_max, out=None)[source]
# 其中a是一个数组，后面两个参数分别表示最小和最大值
result_val = np.clip(result_val, 0, 255)
img_arr = np.asarray(result_val, np.uint8)
img = Image.fromarray(img_arr)
img.save(result_img_path)