class AnoGAN:
"""AnoGAN Class
"""
def __init__(self, opt):
# super(AnoGAN, self).__init__(opt, dataloader)
# Initalize variables.
self.opt = opt
self.niter = self.opt.niter
self.start_iter = 0
self.netd_niter = 5
self.test_iter = 100
self.lr = self.opt.lr
self.batchsize = {'train': self.opt.batchsize, 'test': 1}
self.pretrained = False
self.phase = 'train'
self.outf = self.opt.experiment_group
self.algorithm = 'wgan'
# LOAD DATA SET
self.dataloader = {
'train':
provider('train',
opt.category,
batch_size=self.batchsize['train'],
num_workers=4),
'test':
provider('test',
opt.category,
batch_size=self.batchsize['test'],
num_workers=4)
}
self.trn_dir = os.path.join(self.outf, self.opt.experiment_name,
'train')
self.tst_dir = os.path.join(self.outf, self.opt.experiment_name,
'test')
self.test_img_dir = os.path.join(self.outf, self.opt.experiment_name,
'test', 'images')
if not os.path.isdir(self.test_img_dir):
os.makedirs(self.test_img_dir)
self.best_test_dir = os.path.join(self.outf, self.opt.experiment_name,
'test', 'best_images')
if not os.path.isdir(self.best_test_dir):
os.makedirs(self.best_test_dir)
self.weight_dir = os.path.join(self.trn_dir, 'weights')
if not os.path.exists(self.weight_dir): os.makedirs(self.weight_dir)
# -- Misc attributes
self.epoch = 0
self.l_con = l1_loss
self.l_enc = l2_loss
##
# Create and initialize networks.
self.netg = NetG().cuda()
self.netd = NetD().cuda()
# Setup optimizer
self.optimizer_d = optim.RMSprop(self.netd.parameters(), lr=self.lr)
self.optimizer_g = optim.Adam(self.netg.parameters(), lr=self.lr)
##
self.weight_path = os.path.join(self.outf, self.opt.experiment_name,
'train', 'weights')
if os.path.exists(self.weight_path) and len(
os.listdir(self.weight_path)) == 2:
print("Loading pre-trained networks...\n")
self.netg.load_state_dict(
torch.load(os.path.join(self.weight_path,
'netG.pth'))['state_dict'])
self.netd.load_state_dict(
torch.load(os.path.join(self.weight_path,
'netD.pth'))['state_dict'])
self.optimizer_g.load_state_dict(
torch.load(os.path.join(self.weight_path,
'netG.pth'))['optimizer'])
self.optimizer_d.load_state_dict(
torch.load(os.path.join(self.weight_path,
'netD.pth'))['optimizer'])
self.start_iter = torch.load(
os.path.join(self.weight_path, 'netG.pth'))['epoch']
##
def start(self):
""" Train the model
"""
##
# TRAIN
# self.total_steps = 0
best_criterion = -1 #float('inf')
best_auc = -1
# Train for niter epochs.
# print(">> Training model %s." % self.name)
for self.epoch in range(self.start_iter, self.niter):
# Train for one epoch
mean_wass = self.train()
(auc, res, best_rec, best_threshold), res_total = self.test()
message = ''
# message += 'criterion: (%.3f+%.3f)/2=%.3f ' % (best_rec[0], best_rec[1], res)
# message += 'best threshold: %.3f ' % best_threshold
message += 'Wasserstein Distance:%.3d ' % mean_wass
message += 'AUC: %.3f ' % auc
print(message)
torch.save(
{
'epoch': self.epoch + 1,
'state_dict': self.netg.state_dict(),
'optimizer': self.optimizer_g.state_dict()
}, '%s/netG.pth' % (self.weight_dir))
torch.save(
{
'epoch': self.epoch + 1,
'state_dict': self.netd.state_dict(),
'optimizer': self.optimizer_d.state_dict()
}, '%s/netD.pth' % (self.weight_dir))
if auc > best_auc:
best_auc = auc
new_message = "******** New optimal found, saving state ********"
message = message + '\n' + new_message
print(new_message)
for img in os.listdir(self.best_test_dir):
os.remove(os.path.join(self.best_test_dir, img))
for img in os.listdir(self.test_img_dir):
shutil.copyfile(os.path.join(self.test_img_dir, img),
os.path.join(self.best_test_dir, img))
shutil.copyfile('%s/netG.pth' % (self.weight_dir),
'%s/netg_best.pth' % (self.weight_dir))
log_name = os.path.join(self.outf, self.opt.experiment_name,
'loss_log.txt')
message = 'Epoch%3d:' % self.epoch + ' ' + message
with open(log_name, "a") as log_file:
if self.epoch == 0:
log_file.write('\n\n')
log_file.write('%s\n' % message)
print(">> Training %s Done..." % self.opt.experiment_name)
##
def train(self):
""" Train the model for one epoch.
"""
print("\n>>> Epoch %d/%d, Running " % (self.epoch + 1, self.niter) +
self.opt.experiment_name)
self.netg.train()
self.netd.train()
# for p in self.netg.parameters(): p.requires_grad = True
mean_wass = 0
tk0 = tqdm(self.dataloader['train'],
total=len(self.dataloader['train']))
for i, itr in enumerate(tk0):
input, _ = itr
input = input.cuda()
wasserstein_d = None
# if self.algorithm == 'wgan':
# train NetD
for _ in range(self.netd_niter):
# for p in self.netd.parameters(): p.requires_grad = True
self.optimizer_d.zero_grad()
# forward_g
latent_i = torch.rand(self.batchsize['train'], 64, 1, 1).cuda()
fake = self.netg(latent_i)
# forward_d
_, pred_real = self.netd(input)
_, pred_fake = self.netd(fake) # .detach() TODO
# Backward-pass
wasserstein_d = (pred_fake.mean() - pred_real.mean()) * 1
wasserstein_d.backward()
self.optimizer_d.step()
for p in self.netd.parameters():
p.data.clamp_(-0.01, 0.01) #<<<<<<<
# train netg
# for p in self.netd.parameters(): p.requires_grad = False
self.optimizer_g.zero_grad()
noise = torch.rand(self.batchsize['train'], 64, 1, 1).cuda()
fake = self.netg(noise)
_, pred_fake = self.netd(fake)
err_g_d = -pred_fake.mean() # negative
err_g_d.backward()
self.optimizer_g.step()
errors = {
'loss_netD': wasserstein_d.item(),
'loss_netG': round(err_g_d.item(), 3),
}
mean_wass += wasserstein_d.item()
tk0.set_postfix(errors)
if i % 50 == 0:
img_dir = os.path.join(self.outf, self.opt.experiment_name,
'train', 'images')
if not os.path.isdir(img_dir):
os.makedirs(img_dir)
self.save_image_cv2(input.data, '%s/reals.png' % img_dir)
self.save_image_cv2(fake.data,
'%s/fakes%03d.png' % (img_dir, i))
mean_wass /= len(self.dataloader['train'])
return mean_wass
##
def test(self):
""" Test AnoGAN model.
Args:
dataloader ([type]): Dataloader for the test set
Raises:
IOError: Model weights not found.
"""
self.netg.eval()
self.netd.eval()
# for p in self.netg.parameters(): p.requires_grad = False
# for p in self.netd.parameters(): p.requires_grad = False
for img in os.listdir(self.test_img_dir):
os.remove(os.path.join(self.test_img_dir, img))
self.phase = 'test'
meter = Meter_AnoGAN()
tk1 = tqdm(self.dataloader['test'], total=len(self.dataloader['test']))
for i, itr in enumerate(tk1):
input, target = itr
input = input.cuda()
latent_i = torch.rand(self.batchsize['test'], 64, 1, 1).cuda()
latent_i.requires_grad = True
optimizer_latent = optim.Adam([latent_i], lr=self.lr)
test_loss = None
for _ in range(self.test_iter):
optimizer_latent.zero_grad()
fake = self.netg(latent_i)
residual_loss = self.l_con(input, fake)
latent_o, _ = self.netd(fake)
discrimination_loss = self.l_enc(latent_i, latent_o)
alpha = 0.1
test_loss = (
1 - alpha) * residual_loss + alpha * discrimination_loss
test_loss.backward()
optimizer_latent.step()
abnormal_score = test_loss
meter.update(abnormal_score, target) #<<<TODO
# Save test images.
combine = torch.cat([input.cpu(), fake.cpu()], dim=0)
self.save_image_cv2(combine,
'%s/%05d.jpg' % (self.test_img_dir, i + 1))
criterion, res_total = meter.get_metrics()
# rename images
for i, res in enumerate(res_total):
os.rename('%s/%05d.jpg' % (self.test_img_dir, i + 1),
'%s/%05d_%s.jpg' % (self.test_img_dir, i + 1, res))
return criterion, res_total
@staticmethod
def save_image_cv2(tensor, filename):
# return
from torchvision.utils import make_grid
# tensor = (tensor + 1) / 2
grid = make_grid(tensor, 8, 2, 0, False, None, False)
ndarray = grid.mul_(255).clamp_(0, 255).permute(1, 2, 0).to(
'cpu', torch.uint8).numpy()
cv2.imwrite(filename, ndarray)
class TACGAN():
def __init__(self, args):
self.lr = args.lr
self.cuda = args.use_cuda
self.batch_size = args.batch_size
self.image_size = args.image_size
self.epochs = args.epochs
self.data_root = args.data_root
self.dataset = args.dataset
self.save_dir = args.save_dir
self.save_prefix = args.save_prefix
self.continue_training = args.continue_training
self.continue_epoch = args.continue_epoch
self.netG_path = args.netg_path
self.netD_path = args.netd_path
self.save_after = args.save_after
self.trainset_loader = None
self.evalset_loader = None
self.num_workers = args.num_workers
self.docvec_size = args.docvec_size
self.n_z = args.n_z # length of the noise vector
self.nl_d = args.nl_d
self.nl_g = args.nl_g
self.nf_g = args.nf_g
self.nf_d = args.nf_d
self.bce_loss = nn.BCELoss()
self.nll_loss = nn.NLLLoss()
self.mse_loss = nn.MSELoss()
self.class_filename = args.class_filename
class_path = os.path.join(self.data_root, self.dataset, self.class_filename)
with open(class_path) as f:
self.num_classes = len([l for l in f])
print(self.num_classes)
self.netD = NetD(n_cls=self.num_classes, n_t=self.nl_d, n_f=self.nf_d, docvec_size=self.docvec_size)
self.netG = NetG(n_z=self.n_z, n_l=self.nl_g, n_c=self.nf_g, n_t=self.docvec_size)
if self.continue_training:
self.loadCheckpoints()
# convert to cuda tensors
if self.cuda and torch.cuda.is_available():
print('CUDA is enabled')
self.netD = nn.DataParallel(self.netD).cuda()
self.netG = nn.DataParallel(self.netG).cuda()
self.bce_loss = self.bce_loss.cuda()
self.nll_loss = self.nll_loss.cuda()
self.mse_loss = self.mse_loss.cuda()
# optimizers for netD and netG
self.optimizerD = optim.Adam(params=self.netD.parameters(), lr=self.lr, betas=(0.5, 0.999))
self.optimizerG = optim.Adam(params=self.netG.parameters(), lr=self.lr, betas=(0.5, 0.999))
# create dir for saving checkpoints and other results if do not exist
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
if not os.path.exists(os.path.join(self.save_dir,'netd_checkpoints')):
os.makedirs(os.path.join(self.save_dir,'netd_checkpoints'))
if not os.path.exists(os.path.join(self.save_dir,'netg_checkpoints')):
os.makedirs(os.path.join(self.save_dir,'netg_checkpoints'))
if not os.path.exists(os.path.join(self.save_dir,'generated_images')):
os.makedirs(os.path.join(self.save_dir,'generated_images'))
# start training process
def train(self):
# write to the log file and print it
log_msg = '********************************************\n'
log_msg += ' Training Parameters\n'
log_msg += 'Dataset:%s\nImage size:%dx%d\n'%(self.dataset, self.image_size, self.image_size)
log_msg += 'Batch size:%d\n'%(self.batch_size)
log_msg += 'Number of epochs:%d\nlr:%f\n'%(self.epochs,self.lr)
log_msg += 'nz:%d\nnl-d:%d\nnl-g:%d\n'%(self.n_z, self.nl_d, self.nl_g)
log_msg += 'nf-g:%d\nnf-d:%d\n'%(self.nf_g, self.nf_d)
log_msg += '********************************************\n\n'
print(log_msg)
with open(os.path.join(self.save_dir, 'training_log.txt'),'a') as log_file:
log_file.write(log_msg)
# load trainset and evalset
imtext_ds = ImTextDataset(data_dir=self.data_root, dataset=self.dataset, train=True, image_size=self.image_size)
self.trainset_loader = DataLoader(dataset=imtext_ds, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers)
print("Dataset loaded successfuly")
# load checkpoints for continuing training
# repeat for the number of epochs
netd_losses = []
netg_losses = []
for epoch in range(self.epochs):
netd_loss, netg_loss = self.trainEpoch(epoch)
netd_losses.append(netd_loss)
netg_losses.append(netg_loss)
self.saveGraph(netd_losses,netg_losses)
#self.evalEpoch(epoch)
self.saveCheckpoints(epoch)
# train epoch
def trainEpoch(self, epoch):
self.netD.train() # set to train mode
self.netG.train() #! set to train mode???
netd_loss_sum = 0
netg_loss_sum = 0
start_time = time()
for i, (images, labels, captions) in enumerate(self.trainset_loader):
batch_size = images.size(0) # !batch size my be different (from self.batch_size) for the last batch
images, labels, captions = Variable(images), Variable(labels), Variable(captions) # !labels should be LongTensor
labels = labels.type(torch.FloatTensor) # convert to FloatTensor (from DoubleTensor)
lbl_real = Variable(torch.ones(batch_size, 1))
lbl_fake = Variable(torch.zeros(batch_size, 1))
noise = Variable(torch.randn(batch_size, self.n_z)) # create random noise
noise.data.normal_(0,1) # normalize the noise
rnd_perm1 = torch.randperm(batch_size) # random permutations for different sets of training tuples
rnd_perm2 = torch.randperm(batch_size)
rnd_perm3 = torch.randperm(batch_size)
rnd_perm4 = torch.randperm(batch_size)
if self.cuda:
images, labels, captions = images.cuda(), labels.cuda(), captions.cuda()
lbl_real, lbl_fake = lbl_real.cuda(), lbl_fake.cuda()
noise = noise.cuda()
rnd_perm1, rnd_perm2, rnd_perm3, rnd_perm4 = rnd_perm1.cuda(), rnd_perm2.cuda(), rnd_perm3.cuda(), rnd_perm4.cuda()
############### Update NetD ###############
self.netD.zero_grad()
# train with wrong image, wrong label, real caption
outD_wrong, outC_wrong = self.netD(images[rnd_perm1], captions[rnd_perm2])
# lossD_wrong = self.bce_loss(outD_wrong, lbl_fake)
lossD_wrong = self.bce_loss(outD_wrong, lbl_fake) + self.mse_loss(outD_wrong, lbl_fake)
lossC_wrong = self.bce_loss(outC_wrong, labels[rnd_perm1])
# train with real image, real label, real caption
outD_real, outC_real = self.netD(images, captions)
#lossD_real = self.bce_loss(outD_real, lbl_real)
lossD_real = self.bce_loss(outD_real, lbl_real) + self.mse_loss(outD_real, lbl_real)
lossC_real = self.bce_loss(outC_real, labels)
# train with fake image, real label, real caption
fake = self.netG(noise, captions)
outD_fake, outC_fake = self.netD(fake.detach(), captions[rnd_perm3])
#lossD_fake = self.bce_loss(outD_fake, lbl_fake)
lossD_fake = self.bce_loss(outD_fake, lbl_fake) + self.mse_loss(outD_fake, lbl_fake)
lossC_fake = self.bce_loss(outC_fake, labels[rnd_perm3])
# backward and forwad for NetD
netD_loss = lossC_wrong+lossC_real+lossC_fake + lossD_wrong+lossD_real+lossD_fake
netD_loss.backward()
self.optimizerD.step()
########## Update NetG ##########
# train with fake data
self.netG.zero_grad()
noise.data.normal_(0,1) # normalize the noise vector
fake = self.netG(noise, captions[rnd_perm4])
d_fake, c_fake = self.netD(fake, captions[rnd_perm4])
#lossD_fake_G = self.bce_loss(d_fake, lbl_real)
lossD_fake_G = self.mse_loss(d_fake, lbl_real)
lossC_fake_G = self.bce_loss(c_fake, labels[rnd_perm4])
netG_loss = lossD_fake_G + lossC_fake_G
netG_loss.backward()
self.optimizerG.step()
netd_loss_sum += netD_loss.item()
netg_loss_sum += netG_loss.item()
### print progress info ###
print('Epoch %d/%d, %.2f%% completed. Loss_NetD: %.4f, Loss_NetG: %.4f'
%(epoch, self.epochs,(float(i)/len(self.trainset_loader))*100, netD_loss.item(), netG_loss.item()))
end_time = time()
netd_avg_loss = netd_loss_sum / len(self.trainset_loader)
netg_avg_loss = netg_loss_sum / len(self.trainset_loader)
epoch_time = (end_time-start_time)/60
log_msg = '-------------------------------------------\n'
log_msg += 'Epoch %d took %.2f minutes\n'%(epoch, epoch_time)
log_msg += 'NetD average loss: %.4f, NetG average loss: %.4f\n\n' %(netd_avg_loss, netg_avg_loss)
print(log_msg)
with open(os.path.join(self.save_dir, 'training_log.txt'),'a') as log_file:
log_file.write(log_msg)
return netd_avg_loss, netg_avg_loss
# eval epoch
def evalEpoch(self, epoch):
#self.netD.eval()
#self.netG.eval()
return 0
# draws and saves the loss graph upto the current epoch
def saveGraph(self, netd_losses, netg_losses):
plt.plot(netd_losses, color='red', label='NetD Loss')
plt.plot(netg_losses, color='blue', label='NetG Loss')
plt.xlabel('epoch')
plt.ylabel('error')
plt.legend(loc='best')
plt.savefig(os.path.join(self.save_dir,'loss_graph.png'))
plt.close()
# save after each epoch
def saveCheckpoints(self, epoch):
if epoch%self.save_after==0:
name_netD = "netd_checkpoints/netD_" + self.save_prefix + "_epoch_" + str(epoch) + ".pth"
name_netG = "netg_checkpoints/netG_" + self.save_prefix + "_epoch_" + str(epoch) + ".pth"
torch.save(self.netD.module.state_dict(), os.path.join(self.save_dir, name_netD))
torch.save(self.netG.module.state_dict(), os.path.join(self.save_dir, name_netG))
print("Checkpoints for epoch %d saved successfuly" %(epoch))
# SAVE: data parallel model => add .module
# LOAD: create model and load checkpoints(not add .module) and wrap nn.DataParallel
# this is for fitting prefix
# load checkpoints to continue training
def loadCheckpoints(self):
name_netD = "netd_checkpoints/netD_" + self.save_prefix + "_epoch_" + str(self.continue_epoch) + ".pth"
name_netG = "netg_checkpoints/netG_" + self.save_prefix + "_epoch_" + str(self.continue_epoch) + ".pth"
self.netG.load_state_dict(torch.load(os.path.join(self.save_dir, name_netG)))
self.netD.load_state_dict(torch.load(os.path.join(self.save_dir, name_netD)))
print("Checkpoints loaded successfuly")
all_mods = itertools.chain(all_mods, [
list(netg.children())[0].children(),
list(netd.children())[0].children(),
list(neta.children())[0].children()
])
for mod in all_mods:
if isinstance(mod, nn.Conv2d) or isinstance(mod, nn.ConvTranspose2d):
init.normal_(mod.weight, 0.0, 0.02)
elif isinstance(mod, nn.BatchNorm2d):
init.normal_(mod.weight, 1.0, 0.02)
init.constant_(mod.bias, 0.0)
netg = netg.to(device)
netd = netd.to(device)
neta = neta.to(device)
netg.train()
netd.train()
neta.train()
dataset = CASIABDataset(data_dir='../data/GEI_CASIA_B/gei/')
iteration = 0
lr = 0.0002
real_label = 1
fake_label = 0
fineSize = 64
label = th.zeros((128, 1), requires_grad=False).to(device)
optimG = optim.Adam(netg.parameters(), lr=lr/2)
optimD = optim.Adam(netd.parameters(), lr=lr/3)
optimA = optim.Adam(neta.parameters(), lr=lr/3)
print('Training starts')
def train_network():
init_epoch = 0
best_f1 = 0
total_steps = 0
train_dir = ct.TRAIN_TXT
val_dir = ct.VAL_TXT
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.benchmark = True
train_data = OSCD_TRAIN(train_dir)
train_dataloader = DataLoader(train_data,
batch_size=ct.BATCH_SIZE,
shuffle=True)
val_data = OSCD_TEST(val_dir)
val_dataloader = DataLoader(val_data, batch_size=1, shuffle=False)
netg = NetG(ct.ISIZE, ct.NC * 2, ct.NZ, ct.NDF,
ct.EXTRALAYERS).to(device=device)
netd = NetD(ct.ISIZE, ct.GT_C, 1, ct.NGF, ct.EXTRALAYERS).to(device=device)
netg.apply(weights_init)
netd.apply(weights_init)
if ct.RESUME:
assert os.path.exists(os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth')) \
and os.path.exists(os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth')), \
'There is not found any saved weights'
print("\nLoading pre-trained networks.")
init_epoch = torch.load(
os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth'))['epoch']
netg.load_state_dict(
torch.load(os.path.join(ct.WEIGHTS_SAVE_DIR,
'current_netG.pth'))['model_state_dict'])
netd.load_state_dict(
torch.load(os.path.join(ct.WEIGHTS_SAVE_DIR,
'current_netD.pth'))['model_state_dict'])
with open(os.path.join(ct.OUTPUTS_DIR, 'f1_score.txt')) as f:
lines = f.readlines()
best_f1 = float(lines[-2].strip().split(':')[-1])
print("\tDone.\n")
l_adv = l2_loss
l_con = nn.L1Loss()
l_enc = l2_loss
l_bce = nn.BCELoss()
l_cos = cos_loss
dice = DiceLoss()
optimizer_d = optim.Adam(netd.parameters(), lr=ct.LR, betas=(0.5, 0.999))
optimizer_g = optim.Adam(netg.parameters(), lr=ct.LR, betas=(0.5, 0.999))
start_time = time.time()
for epoch in range(init_epoch + 1, ct.EPOCH):
loss_g = []
loss_d = []
netg.train()
netd.train()
epoch_iter = 0
for i, data in enumerate(train_dataloader):
INPUT_SIZE = [ct.ISIZE, ct.ISIZE]
x1, x2, gt = data
x1 = x1.to(device, dtype=torch.float)
x2 = x2.to(device, dtype=torch.float)
gt = gt.to(device, dtype=torch.float)
gt = gt[:, 0, :, :].unsqueeze(1)
x = torch.cat((x1, x2), 1)
epoch_iter += ct.BATCH_SIZE
total_steps += ct.BATCH_SIZE
real_label = torch.ones(size=(x1.shape[0], ),
dtype=torch.float32,
device=device)
fake_label = torch.zeros(size=(x1.shape[0], ),
dtype=torch.float32,
device=device)
#forward
fake = netg(x)
pred_real = netd(gt)
pred_fake = netd(fake).detach()
err_d_fake = l_bce(pred_fake, fake_label)
err_g = l_con(fake, gt)
err_g_total = ct.G_WEIGHT * err_g + ct.D_WEIGHT * err_d_fake
pred_fake_ = netd(fake.detach())
err_d_real = l_bce(pred_real, real_label)
err_d_fake_ = l_bce(pred_fake_, fake_label)
err_d_total = (err_d_real + err_d_fake_) * 0.5
#backward
optimizer_g.zero_grad()
err_g_total.backward(retain_graph=True)
optimizer_g.step()
optimizer_d.zero_grad()
err_d_total.backward()
optimizer_d.step()
errors = utils.get_errors(err_d_total, err_g_total)
loss_g.append(err_g_total.item())
loss_d.append(err_d_total.item())
counter_ratio = float(epoch_iter) / len(train_dataloader.dataset)
if (i % ct.DISPOLAY_STEP == 0 and i > 0):
print(
'epoch:', epoch, 'iteration:', i,
' G|D loss is {}|{}'.format(np.mean(loss_g[-51:]),
np.mean(loss_d[-51:])))
if ct.DISPLAY:
utils.plot_current_errors(epoch, counter_ratio, errors,
vis)
utils.display_current_images(gt.data, fake.data, vis)
utils.save_current_images(epoch, gt.data, fake.data, ct.IM_SAVE_DIR,
'training_output_images')
with open(os.path.join(ct.OUTPUTS_DIR, 'train_loss.txt'), 'a') as f:
f.write(
'after %s epoch, loss is %g,loss1 is %g,loss2 is %g,loss3 is %g'
% (epoch, np.mean(loss_g), np.mean(loss_d), np.mean(loss_g),
np.mean(loss_d)))
f.write('\n')
if not os.path.exists(ct.WEIGHTS_SAVE_DIR):
os.makedirs(ct.WEIGHTS_SAVE_DIR)
utils.save_weights(epoch, netg, optimizer_g, ct.WEIGHTS_SAVE_DIR,
'netG')
utils.save_weights(epoch, netd, optimizer_d, ct.WEIGHTS_SAVE_DIR,
'netD')
duration = time.time() - start_time
print('training duration is %g' % duration)
#val phase
print('Validating.................')
pretrained_dict = torch.load(
os.path.join(ct.WEIGHTS_SAVE_DIR,
'current_netG.pth'))['model_state_dict']
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = NetG(ct.ISIZE, ct.NC * 2, ct.NZ, ct.NDF,
ct.EXTRALAYERS).to(device=device)
net.load_state_dict(pretrained_dict, False)
with net.eval() and torch.no_grad():
TP = 0
FN = 0
FP = 0
TN = 0
for k, data in enumerate(val_dataloader):
x1, x2, label = data
x1 = x1.to(device, dtype=torch.float)
x2 = x2.to(device, dtype=torch.float)
label = label.to(device, dtype=torch.float)
label = label[:, 0, :, :].unsqueeze(1)
x = torch.cat((x1, x2), 1)
time_i = time.time()
v_fake = net(x)
tp, fp, tn, fn = eva.f1(v_fake, label)
TP += tp
FN += fn
TN += tn
FP += fp
precision = TP / (TP + FP + 1e-8)
oa = (TP + TN) / (TP + FN + TN + FP + 1e-8)
recall = TP / (TP + FN + 1e-8)
f1 = 2 * precision * recall / (precision + recall + 1e-8)
if not os.path.exists(ct.BEST_WEIGHT_SAVE_DIR):
os.makedirs(ct.BEST_WEIGHT_SAVE_DIR)
if f1 > best_f1:
best_f1 = f1
shutil.copy(
os.path.join(ct.WEIGHTS_SAVE_DIR, 'current_netG.pth'),
os.path.join(ct.BEST_WEIGHT_SAVE_DIR, 'netG.pth'))
print('current F1: {}'.format(f1))
print('best f1: {}'.format(best_f1))
with open(os.path.join(ct.OUTPUTS_DIR, 'f1_score.txt'), 'a') as f:
f.write('current epoch:{},current f1:{},best f1:{}'.format(
epoch, f1, best_f1))
f.write('\n')
