class Visualizer(object):
def __init__(self, input_topic, output_topic, resize_width, resize_height,
model_path, force_cpu):
self.bridge = CvBridge()
self.graph = UNet([3, resize_width, resize_height], 3)
self.graph.load_state_dict(torch.load(model_path))
self.force_cpu = force_cpu and torch.cuda.is_available()
self.resize_width, self.resize_height = resize_width, resize_height
if not self.force_cpu:
self.graph.cuda()
self.graph.eval()
self.to_tensor = transforms.Compose([transforms.ToTensor()])
self.publisher = rospy.Publisher(output_topic, ImMsg, queue_size=1)
self.raw_subscriber = rospy.Subscriber(input_topic,
CompressedImage,
self.image_cb,
queue_size=1,
buff_size=10**8)
def convert_to_tensor(self, image):
np_arr = np.fromstring(image.data, np.uint8)
image_np = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
image_np = cv2.resize(image_np,
dsize=(self.resize_width, self.resize_height))
img_to_tensor = PIL.Image.fromarray(image_np)
img_tensor = self.to_tensor(img_to_tensor)
if not self.force_cpu:
return Variable(img_tensor.unsqueeze(0)).cuda()
else:
return Variable(img_tensor.unsqueeze(0))
def image_cb(self, image):
img_tensor = self.convert_to_tensor(image)
# Inference
output = self.graph(img_tensor)
output_data = output.cpu().data.numpy()[0][0]
# # Convert from 32fc1 (0 - 1) to 8uc1 (0 - 255)
cv_output = np.uint8(255 * output_data)
cv_output = cv2.applyColorMap(cv_output, cv2.COLORMAP_JET)
# Convert to ROS message to publish
msg_out = self.bridge.cv2_to_imgmsg(cv_output, 'bgr8')
msg_out.header.stamp = image.header.stamp
self.publisher.publish(msg_out)
def load_finetuned_model(args, baseline_model):
"""
:param args:
:param baseline_model:
:return:
"""
# augment_net = Net(0, 0.0, 32, 3, 0.0, num_classes=32**2 * 3, do_res=True)
augment_net = UNet(in_channels=3, n_classes=3, depth=1, wf=2, padding=True, batch_norm=False,
do_noise_channel=True,
up_mode='upsample', use_identity_residual=True) # TODO(PV): Initialize UNet properly
# TODO (JON): DEPTH 1 WORKED WELL. Changed upconv to upsample. Use a wf of 2.
# This ResNet outputs scalar weights to be applied element-wise to the per-example losses
from models.simple_models import CNN, Net
imsize, in_channel, num_classes = 32, 3, 10
reweighting_net = Net(0, 0.0, imsize, in_channel, 0.0, num_classes=1)
#resnet_cifar.resnet20(num_classes=1)
if args.load_finetune_checkpoint:
checkpoint = torch.load(args.load_finetune_checkpoint)
baseline_model.load_state_dict(checkpoint['elementary_model_state_dict'])
augment_net.load_state_dict(checkpoint['augment_model_state_dict'])
try:
reweighting_net.load_state_dict(checkpoint['reweighting_model_state_dict'])
except KeyError:
pass
augment_net, reweighting_net, baseline_model = augment_net.cuda(), reweighting_net.cuda(), baseline_model.cuda()
augment_net.train(), reweighting_net.train(), baseline_model.train()
return augment_net, reweighting_net, baseline_model
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
model = UNet(3, n_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.cuda(gpu0)
model.train()
testloader = data.DataLoader(REFUGE(False,
domain='REFUGE_TEST',
is_transform=True),
batch_size=args.batch_size,
shuffle=False,
pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(460, 460),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(460, 460), mode='bilinear')
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, label, _, _, name = batch
if args.model == 'Unet':
_, _, _, _, output2 = model(
Variable(image, volatile=True).cuda(gpu0))
output = interp(output2).cpu().data.numpy()
for idx, one_name in enumerate(name):
pred = output[idx]
pred = pred.transpose(1, 2, 0)
pred = np.asarray(np.argmax(pred, axis=2), dtype=np.uint8)
output_col = colorize_mask(pred)
if is_polar:
# plt.imshow(output_col)
# plt.show()
output_col = np.array(output_col)
output_col[output_col == 0] = 0
output_col[output_col == 1] = 128
output_col[output_col == 2] = 255
# plt.imshow(output_col)
# plt.show()
output_col = cv2.linearPolar(
rotate(output_col, 90),
(args.ROI_size / 2, args.ROI_size / 2), args.ROI_size / 2,
cv2.WARP_FILL_OUTLIERS + cv2.WARP_INVERSE_MAP)
# plt.imshow(output_col)
# plt.show()
output_col = np.array(output_col * 255, dtype=np.uint8)
output_col[output_col > 200] = 210
output_col[output_col == 0] = 255
output_col[output_col == 210] = 0
output_col[(output_col > 0) & (output_col < 255)] = 128
output_col = Image.fromarray(output_col)
# plt.imshow(output_col)
# plt.show()
one_name = one_name.split('/')[-1]
if len(one_name.split('_')) > 0:
one_name = one_name[:-4]
#pred.save('%s/%s.bmp' % (args.save, one_name))
output_col = output_col.convert('L')
print(output_col.size)
output_col.save('%s/%s.bmp' % (args.save, one_name.split('.')[0]))
shuffle=False,
num_workers=4)
if args.model == "unet":
model = UNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS)
elif args.model == "segnet":
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_hannels=NUM_OUTPUT_CHANNELS)
else:
model = PSPNet(layers=50, bins=(1, 2, 3, 6), dropout=0.1, classes=NUM_OUTPUT_CHANNELS, use_ppm=True, pretrained=True)
class_weights = 1.0/train_dataset.get_class_probability()
criterion = torch.nn.CrossEntropyLoss(weight=class_weights)
if CUDA:
model = model.cuda(GPU_ID)
class_weights = class_weights.cuda(GPU_ID)
criterion = criterion.cuda(GPU_ID)
if args.checkpoint:
model.load_state_dict(torch.load(args.checkpoint))
optimizer = torch.optim.Adam(model.parameters(),
lr=LEARNING_RATE)
train()
class Noise2Noise(object):
"""Implementation of Noise2Noise from Lehtinen et al. (2018)."""
def __init__(self, params, trainable):
"""Initializes model."""
self.p = params
self.trainable = trainable
self._compile() #初始化模型
def _compile(self):
"""
Compiles model (architecture, loss function, optimizers, etc.).
初始化 网络、损失函数、优化器等
"""
print('Noise2Noise: Learning Image Restoration without Clean Data (Lethinen et al., 2018)')
# Model (3x3=9 channels for Monte Carlo since it uses 3 HDR buffers) 已删除蒙特卡洛相关代码
if self.p.noise_type == 'mc':
self.is_mc = True
self.model = UNet(in_channels=9)
else:
self.is_mc = False
self.model = UNet(in_channels=3)
# Set optimizer and loss, if in training mode
# 如果 为训练，则初始化优化器和损失
if self.trainable:
self.optim = Adam(self.model.parameters(),
lr=self.p.learning_rate,
betas=self.p.adam[:2],
eps=self.p.adam[2])
# Learning rate adjustment
self.scheduler = lr_scheduler.ReduceLROnPlateau(self.optim,
patience=self.p.nb_epochs/4, factor=0.5, verbose=True)
# Loss function
if self.p.loss == 'hdr':
assert self.is_mc, 'Using HDR loss on non Monte Carlo images'
self.loss = HDRLoss()
elif self.p.loss == 'l2':
self.loss = nn.MSELoss()
else:
self.loss = nn.L1Loss()
# CUDA support
self.use_cuda = torch.cuda.is_available() and self.p.cuda
if self.use_cuda:
self.model = self.model.cuda()
if self.trainable:
self.loss = self.loss.cuda()
def _print_params(self):
"""Formats parameters to print when training."""
print('Training parameters: ')
self.p.cuda = self.use_cuda
param_dict = vars(self.p)
pretty = lambda x: x.replace('_', ' ').capitalize()
print('\n'.join(' {} = {}'.format(pretty(k), str(v)) for k, v in param_dict.items()))
print()
def save_model(self, epoch, stats, first=False):
"""Saves model to files; can be overwritten at every epoch to save disk space."""
# Create directory for model checkpoints, if nonexistent
if first:
if self.p.clean_targets:
ckpt_dir_name = f'{datetime.now():{self.p.noise_type}-clean-%H%M}'
else:
ckpt_dir_name = f'{datetime.now():{self.p.noise_type}-%H%M}'
if self.p.ckpt_overwrite:
if self.p.clean_targets:
ckpt_dir_name = f'{self.p.noise_type}-clean'
else:
ckpt_dir_name = self.p.noise_type
self.ckpt_dir = os.path.join(self.p.ckpt_save_path, ckpt_dir_name)
if not os.path.isdir(self.p.ckpt_save_path):
os.mkdir(self.p.ckpt_save_path)
if not os.path.isdir(self.ckpt_dir):
os.mkdir(self.ckpt_dir)
# Save checkpoint dictionary
if self.p.ckpt_overwrite:
fname_unet = '{}/n2n-{}.pt'.format(self.ckpt_dir, self.p.noise_type)
else:
valid_loss = stats['valid_loss'][epoch]
fname_unet = '{}/n2n-epoch{}-{:>1.5f}.pt'.format(self.ckpt_dir, epoch + 1, valid_loss)
print('Saving checkpoint to: {}\n'.format(fname_unet))
torch.save(self.model.state_dict(), fname_unet)
# Save stats to JSON
fname_dict = '{}/n2n-stats.json'.format(self.ckpt_dir)
with open(fname_dict, 'w') as fp:
json.dump(stats, fp, indent=2)
def load_model(self, ckpt_fname):
"""Loads model from checkpoint file."""
print('Loading checkpoint from: {}'.format(ckpt_fname))
if self.use_cuda:
self.model.load_state_dict(torch.load(ckpt_fname))
else:
self.model.load_state_dict(torch.load(ckpt_fname, map_location='cpu'))
def _on_epoch_end(self, stats, train_loss, epoch, epoch_start, valid_loader):
"""Tracks and saves starts after each epoch."""
# Evaluate model on validation set
print('\rTesting model on validation set... ', end='')
epoch_time = time_elapsed_since(epoch_start)[0]
valid_loss, valid_time, valid_psnr = self.eval(valid_loader)
show_on_epoch_end(epoch_time, valid_time, valid_loss, valid_psnr)
# Decrease learning rate if plateau
self.scheduler.step(valid_loss)
# Save checkpoint
stats['train_loss'].append(train_loss)
stats['valid_loss'].append(valid_loss)
stats['valid_psnr'].append(valid_psnr)
self.save_model(epoch, stats, epoch == 0)
def test(self, test_loader, show=1):
"""Evaluates denoiser on test set."""
self.model.train(False)
source_imgs = []
denoised_imgs = []
clean_imgs = []
# Create directory for denoised images
denoised_dir = os.path.dirname(self.p.data)
save_path = os.path.join(denoised_dir, 'denoised')
if not os.path.isdir(save_path):
os.mkdir(save_path)
for batch_idx, (source, target) in enumerate(test_loader):
# Only do first <show> images
if show == 0 or batch_idx >= show:
break
source_imgs.append(source)
clean_imgs.append(target)
if self.use_cuda:
source = source.cuda()
# Denoise
denoised_img = self.model(source).detach()
denoised_imgs.append(denoised_img)
# Squeeze tensors
source_imgs = [t.squeeze(0) for t in source_imgs]
denoised_imgs = [t.squeeze(0) for t in denoised_imgs]
clean_imgs = [t.squeeze(0) for t in clean_imgs]
# Create montage and save images
print('Saving images and montages to: {}'.format(save_path))
for i in range(len(source_imgs)):
img_name = test_loader.dataset.imgs[i]
create_montage(img_name, self.p.noise_type, save_path, source_imgs[i], denoised_imgs[i], clean_imgs[i], show)
def eval(self, valid_loader):
"""Evaluates denoiser on validation set."""
self.model.train(False)
valid_start = datetime.now()
loss_meter = AvgMeter()
psnr_meter = AvgMeter()
for batch_idx, (source, target) in enumerate(valid_loader):
if self.use_cuda:
source = source.cuda()
target = target.cuda()
# Denoise
source_denoised = self.model(source)
# Update loss
loss = self.loss(source_denoised, target)
loss_meter.update(loss.item())
# Compute PSRN
if self.is_mc:
source_denoised = reinhard_tonemap(source_denoised)
# TODO: Find a way to offload to GPU, and deal with uneven batch sizes
for i in range(self.p.batch_size):
source_denoised = source_denoised.cpu()
target = target.cpu()
psnr_meter.update(psnr(source_denoised[i], target[i]).item())
valid_loss = loss_meter.avg
valid_time = time_elapsed_since(valid_start)[0]
psnr_avg = psnr_meter.avg
return valid_loss, valid_time, psnr_avg
def train(self, train_loader, valid_loader):
"""Trains denoiser on training set."""
self.model.train(True)
self._print_params()
num_batches = len(train_loader)
assert num_batches % self.p.report_interval == 0, 'Report interval must divide total number of batches'
# Dictionaries of tracked stats
stats = {'noise_type': self.p.noise_type,
'noise_param': self.p.noise_param,
'train_loss': [],
'valid_loss': [],
'valid_psnr': []}
# Main training loop
train_start = datetime.now()
for epoch in range(self.p.nb_epochs):
print('EPOCH {:d} / {:d}'.format(epoch + 1, self.p.nb_epochs))
# Some stats trackers
epoch_start = datetime.now()
train_loss_meter = AvgMeter()
loss_meter = AvgMeter()
time_meter = AvgMeter()
# Minibatch SGD
for batch_idx, (source, target) in enumerate(train_loader):
batch_start = datetime.now()
progress_bar(batch_idx, num_batches, self.p.report_interval, loss_meter.val)
if self.use_cuda:
source = source.cuda()
target = target.cuda()
# Denoise image
source_denoised = self.model(source)
loss = self.loss(source_denoised, target)
loss_meter.update(loss.item())
# Zero gradients, perform a backward pass, and update the weights
self.optim.zero_grad()
loss.backward()
self.optim.step()
# Report/update statistics
time_meter.update(time_elapsed_since(batch_start)[1])
if (batch_idx + 1) % self.p.report_interval == 0 and batch_idx:
show_on_report(batch_idx, num_batches, loss_meter.avg, time_meter.avg)
train_loss_meter.update(loss_meter.avg)
loss_meter.reset()
time_meter.reset()
# Epoch end, save and reset tracker
self._on_epoch_end(stats, train_loss_meter.avg, epoch, epoch_start, valid_loader)
train_loss_meter.reset()
train_elapsed = time_elapsed_since(train_start)[0]
print('Training done! Total elapsed time: {}\n'.format(train_elapsed))
remove_alpha=True,
one_hot_mask=2)
test_dataset = TestFromFolder(os.path.join(all_datasets,
'stage1_test/loc.csv'),
transform=T.ToTensor(),
remove_alpha=True)
"""
-----------------
----- Model -----
-----------------
"""
s = UNet(3, 2)
t = UNet(3, 3)
if use_gpu:
s.cuda()
t.cuda()
# lr = 0.001 seems to work WITHOUT PRETRAINING
s_optim = optim.Adam(s.parameters(), lr=0.1)
t_optim = optim.Adam(t.parameters(), lr=0.1)
s_scheduler = torch.optim.lr_scheduler.StepLR(s_optim, step_size=10)
t_scheduler = torch.optim.lr_scheduler.StepLR(t_optim, step_size=10)
gan = GANv2(s=s,
s_optim=s_optim,
s_loss=CrossEntropyLoss2d().cuda(),
s_scheduler=s_scheduler,
g=t,
g_optim=t_optim)
model = UNet(n_channels=3, n_classes=1, height=height, width=width)
else:
height, width = 224, 224
model = eval(args.model, {'torch': torch, 'torchvision': torchvision})
if 'mobilenet_v2' in args.model:
model = torch.nn.Sequential(model.features,
torch.nn.AdaptiveAvgPool2d((1, 1)),
torch.nn.Flatten(start_dim=1),
model.classifier[0], model.classifier[1])
if args.check_diff:
disable_dropout(model)
if args.check_runtime:
graph = Graph.create(model, input_shape=(3, height, width))
model.cuda()
solvert = -1
if len(args.solution_file) > 0:
solver_info, solution = load_solution(args.solution_file)
else:
input_ = torch.randn((bs, 3, height, width)).cuda()
solver_info = SolverInfo(bs=bs, model_name=model_name, mode=mode)
solver_info.extract(
graph, input_,
*list(model.state_dict(keep_vars=True).values()))
solver_model = Model(solver_info, budget, args.solver,
args.ablation)
t0 = time()
solution = solver_model.solve()
solvert = time() - t0
pin_memory=True)
val_loader = DataLoader(scan_dataset_val,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
test_loader = DataLoader(scan_dataset_test,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
pseudo = UNet(1, 1)
net = UNet(1, 1, domain_specific=True)
pseudo.cuda()
net.cuda()
# list optimisers here...
# single optimiser variant 1
optimiser_ps = optim.Adam(pseudo.parameters(), lr=learning_rate)
optimiser_net = optim.Adam(net.parameters(), lr=learning_rate)
print('Project name ', project_name)
print('Learning rate ', learning_rate)
print('Epochs ', epochs)
train_loss = []
train_loss_rec = []
train_loss_seg = []
def main():
"""Create the model and start the training."""
args = get_arguments()
cudnn.enabled = True
n_discriminators = 5
# create teacher & student
student_net = UNet(3, n_classes=args.num_classes)
teacher_net = UNet(3, n_classes=args.num_classes)
student_params = list(student_net.parameters())
# teacher doesn't need gradient as it's just a EMA of the student
teacher_params = list(teacher_net.parameters())
for param in teacher_params:
param.requires_grad = False
student_net.train()
student_net.cuda(args.gpu)
teacher_net.train()
teacher_net.cuda(args.gpu)
cudnn.benchmark = True
unsup_weights = [
args.unsup_weight5, args.unsup_weight6, args.unsup_weight7,
args.unsup_weight8, args.unsup_weight9
]
lambda_adv_tgts = [
args.lambda_adv_tgt5, args.lambda_adv_tgt6, args.lambda_adv_tgt7,
args.lambda_adv_tgt8, args.lambda_adv_tgt9
]
# create a list of discriminators
discriminators = []
for dis_idx in range(n_discriminators):
discriminators.append(FCDiscriminator(num_classes=args.num_classes))
discriminators[dis_idx].train()
discriminators[dis_idx].cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
max_iters = args.num_steps * args.iter_size * args.batch_size
src_set = REFUGE(True,
domain='REFUGE_SRC',
is_transform=True,
augmentations=aug_student,
aug_for_target=aug_teacher,
max_iters=max_iters)
src_loader = data.DataLoader(src_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
src_loader_iter = enumerate(src_loader)
tgt_set = REFUGE(True,
domain='REFUGE_DST',
is_transform=True,
augmentations=aug_student,
aug_for_target=aug_teacher,
max_iters=max_iters)
tgt_loader = data.DataLoader(tgt_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
tgt_loader_iter = enumerate(tgt_loader)
student_optimizer = optim.SGD(student_params,
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
teacher_optimizer = optim_weight_ema.WeightEMA(teacher_params,
student_params,
alpha=args.teacher_alpha)
d_optimizers = []
for idx in range(n_discriminators):
optimizer = optim.Adam(discriminators[idx].parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
d_optimizers.append(optimizer)
calc_bce_loss = torch.nn.BCEWithLogitsLoss()
# labels for adversarial training
source_label, tgt_label = 0, 1
for i_iter in range(args.num_steps):
total_seg_loss = 0
seg_loss_vals = [0] * n_discriminators
adv_tgt_loss_vals = [0] * n_discriminators
d_loss_vals = [0] * n_discriminators
unsup_loss_vals = [0] * n_discriminators
for d_optimizer in d_optimizers:
d_optimizer.zero_grad()
adjust_learning_rate_D(d_optimizer, i_iter, args)
student_optimizer.zero_grad()
adjust_learning_rate(student_optimizer, i_iter, args)
for sub_i in range(args.iter_size):
# ******** Optimize source network with segmentation loss ********
# As we don't change the discriminators, their parameters are fixed
for discriminator in discriminators:
for param in discriminator.parameters():
param.requires_grad = False
_, src_batch = src_loader_iter.__next__()
_, _, src_images, src_labels, _ = src_batch
src_images = Variable(src_images).cuda(args.gpu)
# calculate the segmentation losses
sup_preds = list(student_net(src_images))
seg_losses, total_seg_loss = [], 0
for idx, sup_pred in enumerate(sup_preds):
sup_interp_pred = (sup_pred)
# you also can use dice loss like: dice_loss(src_labels, sup_interp_pred)
seg_loss = Weighted_Jaccard_loss(src_labels, sup_interp_pred,
args.class_weights, args.gpu)
seg_losses.append(seg_loss)
total_seg_loss += seg_loss * unsup_weights[idx]
seg_loss_vals[idx] += seg_loss.item() / args.iter_size
_, tgt_batch = tgt_loader_iter.__next__()
tgt_images0, tgt_lbl0, tgt_images1, tgt_lbl1, _ = tgt_batch
tgt_images0 = Variable(tgt_images0).cuda(args.gpu)
tgt_images1 = Variable(tgt_images1).cuda(args.gpu)
# calculate ensemble losses
stu_unsup_preds = list(student_net(tgt_images1))
tea_unsup_preds = teacher_net(tgt_images0)
total_mse_loss = 0
for idx in range(n_discriminators):
stu_unsup_probs = F.softmax(stu_unsup_preds[idx], dim=-1)
tea_unsup_probs = F.softmax(tea_unsup_preds[idx], dim=-1)
unsup_loss = calc_mse_loss(stu_unsup_probs, tea_unsup_probs,
args.batch_size)
unsup_loss_vals[idx] += unsup_loss.item() / args.iter_size
total_mse_loss += unsup_loss * unsup_weights[idx]
total_mse_loss = total_mse_loss / args.iter_size
# As the requires_grad is set to False in the discriminator, the
# gradients are only accumulated in the generator, the target
# student network is optimized to make the outputs of target domain
# images close to the outputs of source domain images
stu_unsup_preds = list(student_net(tgt_images0))
d_outs, total_adv_loss = [], 0
for idx in range(n_discriminators):
stu_unsup_interp_pred = (stu_unsup_preds[idx])
d_outs.append(discriminators[idx](stu_unsup_interp_pred))
label_size = d_outs[idx].data.size()
labels = torch.FloatTensor(label_size).fill_(source_label)
labels = Variable(labels).cuda(args.gpu)
adv_tgt_loss = calc_bce_loss(d_outs[idx], labels)
total_adv_loss += lambda_adv_tgts[idx] * adv_tgt_loss
adv_tgt_loss_vals[idx] += adv_tgt_loss.item() / args.iter_size
total_adv_loss = total_adv_loss / args.iter_size
# requires_grad is set to True in the discriminator, we only
# accumulate gradients in the discriminators, the discriminators are
# optimized to make true predictions
d_losses = []
for idx in range(n_discriminators):
discriminator = discriminators[idx]
for param in discriminator.parameters():
param.requires_grad = True
sup_preds[idx] = sup_preds[idx].detach()
d_outs[idx] = discriminators[idx](sup_preds[idx])
label_size = d_outs[idx].data.size()
labels = torch.FloatTensor(label_size).fill_(source_label)
labels = Variable(labels).cuda(args.gpu)
d_losses.append(calc_bce_loss(d_outs[idx], labels))
d_losses[idx] = d_losses[idx] / args.iter_size / 2
d_losses[idx].backward()
d_loss_vals[idx] += d_losses[idx].item()
for idx in range(n_discriminators):
stu_unsup_preds[idx] = stu_unsup_preds[idx].detach()
d_outs[idx] = discriminators[idx](stu_unsup_preds[idx])
label_size = d_outs[idx].data.size()
labels = torch.FloatTensor(label_size).fill_(tgt_label)
labels = Variable(labels).cuda(args.gpu)
d_losses[idx] = calc_bce_loss(d_outs[idx], labels)
d_losses[idx] = d_losses[idx] / args.iter_size / 2
d_losses[idx].backward()
d_loss_vals[idx] += d_losses[idx].item()
for d_optimizer in d_optimizers:
d_optimizer.step()
total_loss = total_seg_loss + total_adv_loss + total_mse_loss
total_loss.backward()
student_optimizer.step()
teacher_optimizer.step()
log_str = 'iter = {0:7d}/{1:7d}'.format(i_iter, args.num_steps)
log_str += ', total_seg_loss = {0:.3f} '.format(total_seg_loss)
templ = 'seg_losses = [' + ', '.join(['%.2f'] * len(seg_loss_vals))
log_str += templ % tuple(seg_loss_vals) + '] '
templ = 'ens_losses = [' + ', '.join(['%.5f'] * len(unsup_loss_vals))
log_str += templ % tuple(unsup_loss_vals) + '] '
templ = 'adv_losses = [' + ', '.join(['%.2f'] * len(adv_tgt_loss_vals))
log_str += templ % tuple(adv_tgt_loss_vals) + '] '
templ = 'd_losses = [' + ', '.join(['%.2f'] * len(d_loss_vals))
log_str += templ % tuple(d_loss_vals) + '] '
print(log_str)
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
filename = 'UNet' + str(
args.num_steps_stop) + '_v18_weightedclass.pth'
torch.save(teacher_net.cpu().state_dict(),
os.path.join(args.snapshot_dir, filename))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
filename = 'UNet' + str(i_iter) + '_v18_weightedclass.pth'
torch.save(teacher_net.cpu().state_dict(),
os.path.join(args.snapshot_dir, filename))
teacher_net.cuda(args.gpu)
import torch.optim as optim
from models.utils import z, imshow, image_to_tensor, tensor_to_image, crop_image
from models.configs import superresolutionSettings
img_path = "data/superresolution/snail.jpg"
img = Image.open(img_path)
imshow(asarray(img))
img = crop_image(img)
x = image_to_tensor(img)
net = UNet(superresolutionSettings)
if torch.cuda.is_available():
net = net.cuda()
mse = torch.nn.MSELoss()
optimizer = optim.Adam(net.parameters(), lr=0.01)
# Num of iters for training
num_iters = 2000
# Num of iters when to save image
save_frequency = 100
z0 = z(shape=(img.height, img.width), channels=32)
for i in range(num_iters):
optimizer.zero_grad()
output = net(z0)
args = parser.parse_args()
dataset_train = MaterialsDataset("PBR_dataset_256/", test = True)
dataset_test = MaterialsDataset("PBR_dataset_256/", test = False)
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=16, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=16, shuffle=True)
if not torch.cuda.is_available():
print("Warning cuda not found, falling back on CPU!")
args.cuda = False
G = UNet(n_channels=3, n_output=3)
D = Discriminator()
if args.cuda:
G.cuda()
D.cuda()
lr=0.001
# run forever
for epoch in range(10000):
lr *= 0.9995
print("Epoch {}, lr: {}".format(epoch,lr) )
G_solver = torch.optim.Adam(model.parameters(),lr=lr )
D_solver = torch.optim.Adam(model.parameters(),lr=lr )
train_loss = list()
tic = time()
# Train Epoch
for batch_idx, batch_item in enumerate(train_loader):
if args.model == 'unet':
model = UNet(3, n_classes=args.num_classes, first_conv_channels=args.unet_channels, batch_norm=unet_batch_norm, dropout_rate=unet_dropout_rate)
else:
raise Exception('Unknown models')
if load_model:
model_path = os.path.join(args.output_dir, 'checkpoint', args.initial_checkpoint)
print_to_log('loading models from file', args.initial_checkpoint, log_file)
model.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage))
# model.load_state_dict(torch.load(model_path))
# torch.load(model_path)
print_to_log('models loaded!', '', log_file)
print("Running models: " + args.model)
cuda = torch.cuda.is_available() and args.use_gpu
if cuda:
model = model.cuda()
else:
model = model.cpu()
print_to_log('gpu', cuda, log_file)
log_file.close()
if not predict:
print("training on train set")
train_set = SemanticSegmentationDataset(args.train_dir,
os.path.join('image_sets/', args.train_set),
image_size, mode='train')
val_set = SemanticSegmentationDataset(args.validation_dir,
os.path.join('image_sets/', args.valid_set),
image_size, mode='valid')
loss = torch.nn.CrossEntropyLoss()
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
test_b_loader_gold = DataLoader(test_b_dataset_gold,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
pseudo = UNet(1, 1)
refine_net = UNet(1, 1)
# UNet(1, 1, norm_layer=nn.BatchNorm2d, affine=True)
net = UNet(1, 1, domain_specific=True)
pseudo.cuda()
refine_net.cuda()
net.cuda()
# list optimisers here...
# single optimiser variant 1
optimiser_ps = optim.Adam(pseudo.parameters(), lr=learning_rate)
optimiser_ref = optim.Adam(refine_net.parameters(), lr=learning_rate * 10)
optimiser_net = optim.Adam(net.parameters(), lr=learning_rate)
scheduler_ps = None # StepLR(optimiser_ps, step_size=1, gamma=0.5)
scheduler_net = None # StepLR(optimiser_net, step_size=5, gamma=0.8)
print('Project name ', project_name)
print('Learning rate ', learning_rate)
'--load',
dest='load',
default=False,
help='load file model')
(options, args) = parser.parse_args()
if (options.model == 1):
net = UNet(3, 1)
if options.load:
net.load_state_dict(torch.load(options.load))
print('Model loaded from {}'.format(options.load))
if options.gpu:
net.cuda()
cudnn.benchmark = True
try:
train_net(net,
options.epochs,
options.batchsize,
options.lr,
gpu=options.gpu)
except KeyboardInterrupt:
torch.save(net.state_dict(), 'INTERRUPTED.pth')
print('Saved interrupt')
try:
sys.exit(0)
except SystemExit:
os._exit(0)
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
test_loader = DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
refine_net = UNet(1,
1,
norm_layer=nn.BatchNorm2d,
affine=True,
track_running_stats=True)
refine_net.cuda()
# list optimisers here...
# single optimiser variant 1
optimiser_ref = optim.Adam(refine_net.parameters(), lr=learning_rate)
print('Project name ', project_name)
train_dices = []
train_losses = []
val_dices = []
val_losses = []
for i in range(epochs):
train_dice, train_loss = train_segmentation_net(refine_net,
num_workers=5,
pin_memory=True)
val_b_loader = DataLoader(val_b_dataset,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
test_b_loader = DataLoader(test_b_dataset,
batch_size=1,
shuffle=False,
num_workers=5,
pin_memory=True)
# net and optimizer
ds_unet = UNet(1, 1, domain_specific=True)
ds_unet.cuda()
labeller = UNet(1, 1)
# import weights here...
labeller_path = './results/unet_sobel_eadan_in/net'
labeller.load_state_dict(
torch.load(os.path.join(labeller_path),
map_location=lambda storage, loc: storage))
labeller.cuda()
optimiser = optim.Adam(ds_unet.parameters(), lr=learning_rate)
print('Project name ', project_name)
train_dices = []
train_losses = []
val_a_dices = []
all_datasets, 'stage1_train_merged/loc.csv'),
transform=T.ToTensor(),
remove_alpha=True)
test_dataset = TestFromFolder(os.path.join(all_datasets,
'stage1_test/loc.csv'),
transform=T.ToTensor(),
remove_alpha=True)
"""
-----------------
----- Model -----
-----------------
"""
generator = UNet(3, 1)
discriminator = Discriminator(4, 1)
generator.cuda()
discriminator.cuda()
# lr = 0.001 seems to work WITHOUT PRETRAINING
g_optim = optim.Adam(generator.parameters(), lr=0.001)
d_optim = optim.Adam(discriminator.parameters(), lr=0.001)
#g_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(g_optim, factor=0.1, verbose=True, patience=5)
#d_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(d_optim, factor=0.1, verbose=True, patience=5)
gan = GAN(
g=generator,
d=discriminator,
g_optim=g_optim,
d_optim=d_optim,
g_loss=nn.MSELoss().cuda(),
d_loss=nn.MSELoss().cuda(),
#g_scheduler=g_scheduler, d_scheduler=d_scheduler
model = UNet(input_channels=NUM_INPUT_CHANNELS,
output_channels=NUM_OUTPUT_CHANNELS)
elif args.model == "segnet":
model = SegNet(input_channels=NUM_INPUT_CHANNELS,
output_hannels=NUM_OUTPUT_CHANNELS)
else:
model = PSPNet(
layers=50,
bins=(1, 2, 3, 6),
dropout=0.1,
classes=NUM_OUTPUT_CHANNELS,
use_ppm=True,
pretrained=True,
)
# class_weights = 1.0 / train_dataset.get_class_probability()
# criterion = torch.nn.CrossEntropyLoss(weight=class_weights)
criterion = torch.nn.CrossEntropyLoss()
if CUDA:
model = model.cuda(device=GPU_ID)
# class_weights = class_weights.cuda(GPU_ID)
criterion = criterion.cuda(device=GPU_ID)
if args.checkpoint:
model.load_state_dict(torch.load(args.checkpoint))
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
train()
def train_eval_model(opts):
# parse model configuration
num_epochs = opts["num_epochs"]
train_batch_size = opts["train_batch_size"]
val_batch_size = opts["eval_batch_size"]
dataset_type = opts["dataset_type"]
opti_mode = opts["optimizer"]
loss_criterion = opts["loss_criterion"]
lr = opts["lr"]
lr_decay = opts["lr_decay"]
wd = opts["weight_decay"]
gpus = opts["gpu_list"].split(',')
os.environ['CUDA_VISIBLE_DEVICE'] = opts["gpu_list"]
train_dir = opts["log_dir"]
train_data_dir = opts["train_data_dir"]
eval_data_dir = opts["eval_data_dir"]
pretrained = opts["pretrained_model"]
resume = opts["resume"]
display_iter = opts["display_iter"]
save_epoch = opts["save_every_epoch"]
show = opts["vis"]
# backup train configs
log_file = os.path.join(train_dir, "log_file.txt")
os.makedirs(train_dir, exist_ok=True)
model_dir = os.path.join(train_dir, "code_backup")
os.makedirs(model_dir, exist_ok=True)
if resume is None and os.path.exists(log_file): os.remove(log_file)
shutil.copy("./models/unet.py", os.path.join(model_dir, "unet.py"))
shutil.copy("./trainer_unet.py", os.path.join(model_dir,
"trainer_unet.py"))
shutil.copy("./datasets/dataset.py", os.path.join(model_dir, "dataset.py"))
ckt_dir = os.path.join(train_dir, "checkpoints")
os.makedirs(ckt_dir, exist_ok=True)
# format printing configs
print("*" * 50)
table_key = []
table_value = []
n = 0
for key, value in opts.items():
table_key.append(key)
table_value.append(str(value))
n += 1
print_table([table_key, ["="] * n, table_value])
# format gpu list
gpu_list = []
for str_id in gpus:
id = int(str_id)
gpu_list.append(id)
# dataloader
print("==> Create dataloader")
dataloaders_dict = {
"train":
er_data_loader(train_data_dir,
train_batch_size,
dataset_type,
is_train=True),
"eval":
er_data_loader(eval_data_dir,
val_batch_size,
dataset_type,
is_train=False)
}
# define parameters of two networks
print("==> Create network")
num_channels = 1
num_classes = 1
model = UNet(num_channels, num_classes)
init_weights(model)
# loss layer
criterion = create_criterion(criterion=loss_criterion)
best_acc = 0.0
start_epoch = 0
# load pretrained model
if pretrained is not None and os.path.isfile(pretrained):
print("==> Train from model '{}'".format(pretrained))
checkpoint_gan = torch.load(pretrained)
model.load_state_dict(checkpoint_gan['model_state_dict'])
print("==> Loaded checkpoint '{}')".format(pretrained))
for param in model.parameters():
param.requires_grad = False
# resume training
elif resume is not None and os.path.isfile(resume):
print("==> Resume from checkpoint '{}'".format(resume))
checkpoint = torch.load(resume)
start_epoch = checkpoint['epoch'] + 1
best_acc = checkpoint['best_acc']
model_dict = model.state_dict()
pretrained_dict = {
k: v
for k, v in checkpoint['model_state_dict'].items()
if k in model_dict and v.size() == model_dict[k].size()
}
model_dict.update(pretrained_dict)
model.load_state_dict(pretrained_dict)
print("==> Loaded checkpoint '{}' (epoch {})".format(
resume, checkpoint['epoch'] + 1))
# train from scratch
else:
print("==> Train from initial or random state.")
# define mutiple-gpu mode
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.cuda()
model = nn.DataParallel(model)
# print learnable parameters
print("==> List learnable parameters")
for name, param in model.named_parameters():
if param.requires_grad == True:
print("\t{}, size {}".format(name, param.size()))
params_to_update = [{'params': model.parameters()}]
# define optimizer
print("==> Create optimizer")
optimizer = create_optimizer(params_to_update,
opti_mode,
lr=lr,
momentum=0.9,
wd=wd)
if resume is not None and os.path.isfile(resume):
optimizer.load_state_dict(checkpoint['optimizer'])
# start training
since = time.time()
# Each epoch has a training and validation phase
print("==> Start training")
total_steps = 0
for epoch in range(start_epoch, num_epochs):
print('-' * 50)
print("==> Epoch {}/{}".format(epoch + 1, num_epochs))
total_steps = train_one_epoch(epoch, total_steps,
dataloaders_dict['train'], model, device,
criterion, optimizer, lr, lr_decay,
display_iter, log_file, show)
epoch_acc, epoch_iou, epoch_f1 = eval_one_epoch(
epoch, dataloaders_dict['eval'], model, device, log_file)
if best_acc < epoch_acc and epoch >= 5:
best_acc = epoch_acc
torch.save(
{
'epoch': epoch,
'model_state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_acc': best_acc
}, os.path.join(ckt_dir, "best.pth"))
if (epoch + 1) % save_epoch == 0 and (epoch + 1) >= 20:
torch.save(
{
'epoch': epoch,
'model_state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
'best_iou': epoch_iou
},
os.path.join(ckt_dir,
"checkpoints_" + str(epoch + 1) + ".pth"))
time_elapsed = time.time() - since
time_message = 'Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60)
print(time_message)
with open(log_file, "a+") as fid:
fid.write('%s\n' % time_message)
print('==> Best val Acc: {:4f}'.format(best_acc))
def train(frame_num,
layer_nums,
input_channels,
output_channels,
discriminator_num_filters,
bn=False,
pretrain=False,
generator_pretrain_path=None,
discriminator_pretrain_path=None):
generator = UNet(n_channels=input_channels,
layer_nums=layer_nums,
output_channel=output_channels,
bn=bn)
discriminator = PixelDiscriminator(output_channels,
discriminator_num_filters,
use_norm=False)
generator = generator.cuda()
discriminator = discriminator.cuda()
flow_network = Network()
flow_network.load_state_dict(torch.load(lite_flow_model_path))
flow_network.cuda().eval()
adversarial_loss = Adversarial_Loss().cuda()
discriminate_loss = Discriminate_Loss().cuda()
gd_loss = Gradient_Loss(alpha, num_channels).cuda()
op_loss = Flow_Loss().cuda()
int_loss = Intensity_Loss(l_num).cuda()
step = 0
if not pretrain:
generator.apply(weights_init_normal)
discriminator.apply(weights_init_normal)
else:
assert (generator_pretrain_path != None
and discriminator_pretrain_path != None)
generator.load_state_dict(torch.load(generator_pretrain_path))
discriminator.load_state_dict(torch.load(discriminator_pretrain_path))
step = int(generator_pretrain_path.split('-')[-1])
print('pretrained model loaded!')
print('initializing the model with Generator-Unet {} layers,'
'PixelDiscriminator with filters {} '.format(
layer_nums, discriminator_num_filters))
optimizer_G = torch.optim.Adam(generator.parameters(), lr=g_lr)
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=d_lr)
writer = SummaryWriter(writer_path)
dataset = img_dataset.ano_pred_Dataset(training_data_folder, frame_num)
dataset_loader = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1,
drop_last=True)
test_dataset = img_dataset.ano_pred_Dataset(testing_data_folder, frame_num)
test_dataloader = DataLoader(dataset=test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1,
drop_last=True)
for epoch in range(epochs):
for (input, _), (test_input, _) in zip(dataset_loader,
test_dataloader):
# generator = generator.train()
# discriminator = discriminator.train()
target = input[:, -1, :, :, :].cuda()
input = input[:, :-1, ]
input_last = input[:, -1, ].cuda()
input = input.view(input.shape[0], -1, input.shape[-2],
input.shape[-1]).cuda()
test_target = test_input[:, -1, ].cuda()
test_input = test_input[:, :-1].view(test_input.shape[0], -1,
test_input.shape[-2],
test_input.shape[-1]).cuda()
#------- update optim_G --------------
G_output = generator(input)
pred_flow_esti_tensor = torch.cat([input_last, G_output], 1)
gt_flow_esti_tensor = torch.cat([input_last, target], 1)
flow_gt = batch_estimate(gt_flow_esti_tensor, flow_network)
flow_pred = batch_estimate(pred_flow_esti_tensor, flow_network)
g_adv_loss = adversarial_loss(discriminator(G_output))
g_op_loss = op_loss(flow_pred, flow_gt)
g_int_loss = int_loss(G_output, target)
g_gd_loss = gd_loss(G_output, target)
g_loss = lam_adv * g_adv_loss + lam_gd * g_gd_loss + lam_op * g_op_loss + lam_int * g_int_loss
optimizer_G.zero_grad()
g_loss.backward()
optimizer_G.step()
train_psnr = psnr_error(G_output, target)
#----------- update optim_D -------
optimizer_D.zero_grad()
d_loss = discriminate_loss(discriminator(target),
discriminator(G_output.detach()))
#d_loss.requires_grad=True
d_loss.backward()
optimizer_D.step()
#----------- cal psnr --------------
test_generator = generator.eval()
test_output = test_generator(test_input)
test_psnr = psnr_error(test_output, test_target).cuda()
if step % 10 == 0:
print("[{}/{}]: g_loss: {} d_loss {}".format(
step, epoch, g_loss, d_loss))
print('\t gd_loss {}, op_loss {}, int_loss {} ,'.format(
g_gd_loss, g_op_loss, g_int_loss))
print('\t train psnr{}，test_psnr {}'.format(
train_psnr, test_psnr))
writer.add_scalar('psnr/train_psnr',
train_psnr,
global_step=step)
writer.add_scalar('psnr/test_psnr',
test_psnr,
global_step=step)
writer.add_scalar('total_loss/g_loss',
g_loss,
global_step=step)
writer.add_scalar('total_loss/d_loss',
d_loss,
global_step=step)
writer.add_scalar('g_loss/adv_loss',
g_adv_loss,
global_step=step)
writer.add_scalar('g_loss/op_loss',
g_op_loss,
global_step=step)
writer.add_scalar('g_loss/int_loss',
g_int_loss,
global_step=step)
writer.add_scalar('g_loss/gd_loss',
g_gd_loss,
global_step=step)
writer.add_image('image/train_target',
target[0],
global_step=step)
writer.add_image('image/train_output',
G_output[0],
global_step=step)
writer.add_image('image/test_target',
test_target[0],
global_step=step)
writer.add_image('image/test_output',
test_output[0],
global_step=step)
step += 1
if step % 500 == 0:
utils.saver(generator.state_dict(),
model_generator_save_path,
step,
max_to_save=10)
utils.saver(discriminator.state_dict(),
model_discriminator_save_path,
step,
max_to_save=10)
if step >= 2000:
print('==== begin evaluate the model of {} ===='.format(
model_generator_save_path + '-' + str(step)))
auc = evaluate(frame_num=5,
layer_nums=4,
input_channels=12,
output_channels=3,
model_path=model_generator_save_path + '-' +
str(step),
evaluate_name='compute_auc')
writer.add_scalar('results/auc', auc, global_step=step)
scan_dataset, _ = random_split(train_scan_dataset, [num_train, len(train_scan_dataset) - num_train])
scan_dataset_test, scan_dataset_val, _ = random_split(test_scan_dataset,
[num_val, num_test, len(test_scan_dataset) - (num_val + num_test)])
train_loader = DataLoader(scan_dataset, batch_size=batch_size, num_workers=5, pin_memory=True)
val_loader = DataLoader(scan_dataset_val, batch_size=1, shuffle=False, num_workers=5, pin_memory=True)
test_loader = DataLoader(scan_dataset_test, batch_size=1, shuffle=False, num_workers=5, pin_memory=True)
pseudo = UNet(1, 1)
# this one works best!
net = UNetDense(1, 1, norm_layer=nn.BatchNorm2d, track_running_stats=False, affine=True)
# net = UNetWavelet(1, 1)
pseudo.cuda()
net.cuda()
# list optimisers here...
# single optimiser variant 1
optimiser_ps = optim.Adam(pseudo.parameters(), lr=learning_rate)
optimiser_net = optim.Adam(net.parameters(), lr=learning_rate)
print('Project name ', project_name)
print('Learning rate ', learning_rate)
print('Epochs ', epochs)
train_loss = []
train_loss_rec = []
train_loss_seg = []
