def __init__(self):
# Device
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
# Model
self.net = net.MyNet().to(self.device)
if not os.path.exists(cfg.PATH) or not os.path.exists(cfg.IMG):
os.makedirs(cfg.PATH)
os.makedirs(cfg.IMG)
if not os.path.exists(cfg.MODEL):
print("Initing ... ...")
self.net.apply(utils.weights_init)
else:
print("Loading ... ...")
self.net.load_state_dict(torch.load(cfg.MODEL))
# Data
self.train = dataset.TRAIN
self.test = dataset.TEST
# Loss
self.centerloss = net.CenterLoss(cfg.CLS_NUM,
cfg.FEATURE_NUM).to(self.device)
if os.path.exists(cfg.CLOSS):
self.centerloss.load_state_dict(torch.load(cfg.CLOSS))
self.clsloss = nn.CrossEntropyLoss()
# Optimize
self.opt = optim.Adam(self.net.parameters())
self.opt_centerloss = optim.Adam(self.centerloss.parameters())
def __init__(self):
# Device
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
# Model
self.net = net.MyNet().to(self.device)
if not os.path.exists(cfg.PATH) or not os.path.exists(cfg.IMG):
os.makedirs(cfg.PATH)
os.makedirs(cfg.IMG)
if not os.path.exists(cfg.MODEL):
print("Initing ... ...")
self.net.apply(utils.weights_init)
else:
print("Loading ... ...")
self.net.load_state_dict(torch.load(cfg.MODEL))
# Data
self.train = dataset.TRAIN
self.test = dataset.TEST
# Optimize
self.opt = optim.Adam(self.net.parameters())
def __init__(self):
# Device
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
# Model
self.net = net.MyNet().to(self.device)
if not os.path.exists(cfg.PARAMS):
if not os.path.exists("./params"):
os.mkdir("./params")
print("Initing ... ...")
self.net.apply(utils.weights_init)
else:
print("Loading ... ...")
self.net.load_state_dict(torch.load(cfg.PARAMS))
# Data
self.train = dataset.TRAIN
self.test = dataset.TEST
# Loss
self.loss = nn.CrossEntropyLoss()
# Optimize
self.opt = optim.Adam(self.net.parameters())
def train(config):
dehaze_net = net.MyNet().cuda()
dehaze_net = nn.DataParallel(dehaze_net, device_ids=[0, 1, 2, 3])
#dehaze_net.apply(weights_init)
dehaze_net.load_state_dict(torch.load('snapshots/dehazer.pth'))
train_dataset = dataloader.dehazing_loader(config.orig_images_path,
config.hazy_images_path)
val_dataset = dataloader.dehazing_loader(config.orig_images_path,
config.hazy_images_path,
mode="val")
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=config.num_workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=config.val_batch_size,
shuffle=True,
num_workers=config.num_workers,
pin_memory=True)
criterion = nn.MSELoss().cuda()
optimizer = torch.optim.Adam(dehaze_net.parameters(),
lr=config.lr,
weight_decay=config.weight_decay)
dehaze_net.train()
for epoch in range(config.num_epochs):
for iteration, (img_orig, img_haze) in enumerate(train_loader):
img_orig = img_orig.cuda()
img_haze = img_haze.cuda()
clean_image = dehaze_net(img_haze)
loss = criterion(clean_image, img_orig)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm(dehaze_net.parameters(),
config.grad_clip_norm)
optimizer.step()
if ((iteration + 1) % config.display_iter) == 0:
print("Loss at iteration", iteration + 1, ":", loss.item())
if ((iteration + 1) % config.snapshot_iter) == 0:
torch.save(
dehaze_net.state_dict(),
config.snapshots_folder + "Epoch" + str(epoch) + '.pth')
# Validation Stage
for iter_val, (img_orig, img_haze) in enumerate(val_loader):
img_orig = img_orig.cuda()
img_haze = img_haze.cuda()
clean_image = dehaze_net(img_haze)
torchvision.utils.save_image(
torch.cat((img_haze, clean_image, img_orig), 0),
config.sample_output_folder + str(iter_val + 1) + ".png")
torch.save(dehaze_net.state_dict(),
config.snapshots_folder + "dehazer.pth")