def test_single(img, config):
model = SRNet()
if config.data_parallel:
model = nn.DataParallel(model)
if config.cuda:
model.cuda()
if config.data_parallel:
load_model(model,
config.model_path,
config.model_name,
mode="parallel")
else:
load_model(model, config.model_path, config.model_name, mode="single")
_img = data_utils.convert_to_torch_tensor(img)
out = model(_img)
out = out.data.cpu().numpy()
if config.hist_eq:
out = image_utils.hist_match(out[0], img)
return out
def test_batch(inp, config):
model = SRNet()
if config.data_parallel:
model = nn.DataParallel(model)
if config.cuda:
model.cuda()
if config.data_parallel:
load_model(model,
config.model_path,
config.model_name,
mode="parallel")
else:
load_model(model, config.model_path, config.model_name, mode="single")
_inp = copy.deepcopy(inp)
_inp = data_utils.convert_to_torch_tensor(_inp)
out = model(_inp)
out = out.data.cpu().numpy()
if config.hist_eq:
out_arr = []
for i in range(out.shape[0]):
out_arr.append(image_utils.hist_match(out[i], inp[i]))
return np.array(out_arr)
def ug2_classifier_loss(y_pred, config):
mapped_output = []
for j in range(len(config.mapping_list)):
mapping = data_utils.convert_to_torch_tensor(np.array(
config.mapping_list[j]),
from_numpy=True,
cuda=False,
dtype="int64")
group = torch.index_select(y_pred, 1, mapping)
mapped_output.append(torch.sum(group, dim=1))
y_pred = torch.stack(mapped_output, dim=1)
def train(config):
model = SRNet()
if config.weights_init:
model.apply(weights_init)
stop_training_flag = False
feat_extractor = feat_ext(config.ext_type, cuda=config.cuda)
if config.discriminator == "classifier":
classifier = Classifier(
pretrained_classifier(config.classifier_type, cuda=config.cuda),
(224, 224))
if config.cuda:
model.cuda()
feat_extractor.cuda()
if config.discriminator == "classifier":
classifier.cuda()
if config.data_parallel:
model = nn.DataParallel(model)
feat_extractor = nn.DataParallel(feat_extractor)
if config.discriminator == "classifier":
classifier = nn.DataParallel(classifier)
if config.resume_training_flag:
if config.data_parallel:
load_model(model,
config.resume_model_path,
config.resume_model_name,
mode="parallel")
else:
load_model(model,
config.resume_model_path,
config.resume_model_name,
mode="single")
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=config.lr)
# scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones = config.lr_step_list, gamma= config.step_decay_param)
loss_fn = nn.MSELoss()
if config.cuda:
loss_fn.cuda()
loss_history = []
epoch_loss_arr = []
for i in range(config.epochs):
loss_arr = []
start = time.time()
if stop_training_flag:
break
for data_file in config.data_files:
dataset = data_utils.DatasetFromFile(
config.data_path,
data_file,
config.img_size,
data_format=config.data_format)
try:
data_loader = DataLoader(dataset,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers)
for batch in data_loader:
x = data_utils.convert_to_torch_tensor(batch["data"],
from_numpy=False,
cuda=config.cuda)
y = data_utils.convert_to_torch_tensor(batch["label"],
from_numpy=False,
cuda=config.cuda)
if config.discriminator == "classifier":
cl = data_utils.convert_to_torch_tensor(
batch["class"], from_numpy=False, cuda=config.cuda)
optimizer.zero_grad()
y_pred = model(x)
y_pred_feat = feat_extractor(y_pred)
feat_loss = loss_fn(y_pred_feat, feat_extractor(y))
total_loss = feat_loss
if config.discriminator == "classifier":
y_pred_imgnet_class = classifier(y_pred)
y_pred_class = ug2_classifier_loss(
y_pred_imgnet_class.cpu(), config)
class_loss = loss_fn(y_pred_class, cl.cpu())
total_loss = feat_loss.cpu() + class_loss
loss_arr.append(total_loss.data[0])
total_loss.backward()
optimizer.step()
except KeyboardInterrupt:
stop_training_flag = True
mean_epoch_loss = np.mean(loss_arr)
epoch_loss_arr.append(mean_epoch_loss)
if i % config.print_step == 0:
print("time: ", time.time() - start, " Error: ", mean_epoch_loss)
if i % config.checkpoint == 0:
if config.save_separate_chkpt:
save_model(model,
optimizer,
path=config.model_path,
filename=config.model_name + "_chkpt" + str(i) +
".pth")
else:
save_model(model,
optimizer,
path=config.model_path,
filename=config.model_name + "_chkpt.pth")
print("saved checkpoint at epoch: ", i)
if exit_training(mean_epoch_loss,
loss_history,
window=config.exit_loss_window,
epsilon=config.loss_epsilon):
stop_training_flag = True
save_model(model,
optimizer,
path=config.model_path,
filename=config.model_name)
print("Model saved as: ", config.model_name)
with open(os.path.join(config.model_path, config.model_name + ".txt"),
'w') as outfile:
json.dump(epoch_loss_arr, outfile)