def __init__(self, weights=False):
super(VGG, self).__init__()
if weights is False:
self.model = models.vgg19(pretrained=True)
else:
self.model = models.vgg19(pretrained=False)
pre = torch.load(weights)
self.model.load_state_dict(pre)
# output features not predictions.
self.vgg19 = self.model.features
for param in self.vgg19.parameters():
param.requires_grad = False
def get_vanilla_vgg_features(cut_idx=-1):
if not os.path.exists('vgg_features.pth'):
os.system(
'wget --no-check-certificate -N https://s3-us-west-2.amazonaws.com/jcjohns-models/vgg19-d01eb7cb.pth')
vgg_weights = torch.load('vgg19-d01eb7cb.pth')
# fix compatibility issues
map = {'classifier.6.weight':u'classifier.7.weight', 'classifier.6.bias':u'classifier.7.bias'}
vgg_weights = OrderedDict([(map[k] if k in map else k,v) for k,v in vgg_weights.iteritems()])
model = models.vgg19()
model.classifier = nn.Sequential(View(), *model.classifier._modules.values())
model.load_state_dict(vgg_weights)
torch.save(model.features, 'vgg_features.pth')
torch.save(model.classifier, 'vgg_classifier.pth')
vgg = torch.load('vgg_features.pth')
if cut_idx > 36:
vgg_classifier = torch.load('vgg_classifier.pth')
vgg = nn.Sequential(*(vgg._modules.values() + vgg_classifier._modules.values()))
vgg.eval()
return vgg
def __init__(self):
super(VGGNet, self).__init__()
"""Select conv1_1 ~ conv5_1 activation maps."""
#self.select = ['0', '5', '10', '19', '28']
self.select = ['0','2','5','7','10','12','14','16','19','21','23','25','28','30','32','34']
self.vgg = models.vgg19(pretrained=True).features
def __init__(self):
super(FeatureExtractor, self).__init__()
vgg19_model = vgg19(pretrained=True)
# Extracts features at the 11th layer
self.feature_extractor = nn.Sequential(*list(vgg19_model.features.children())[:12])
def get_pretrained_net(name):
"""Loads pretrained network"""
if name == 'alexnet_caffe':
if not os.path.exists('alexnet-torch_py3.pth'):
print('Downloading AlexNet')
os.system('wget -O alexnet-torch_py3.pth --no-check-certificate -nc https://box.skoltech.ru/index.php/s/77xSWvrDN0CiQtK/download')
return torch.load('alexnet-torch_py3.pth')
elif name == 'vgg19_caffe':
if not os.path.exists('vgg19-caffe-py3.pth'):
print('Downloading VGG-19')
os.system('wget -O vgg19-caffe-py3.pth --no-check-certificate -nc https://box.skoltech.ru/index.php/s/HPcOFQTjXxbmp4X/download')
vgg = get_vgg19_caffe()
return vgg
elif name == 'vgg16_caffe':
if not os.path.exists('vgg16-caffe-py3.pth'):
print('Downloading VGG-16')
os.system('wget -O vgg16-caffe-py3.pth --no-check-certificate -nc https://box.skoltech.ru/index.php/s/TUZ62HnPKWdxyLr/download')
vgg = get_vgg16_caffe()
return vgg
elif name == 'vgg19_pytorch_modified':
# os.system('wget -O data/feature_inversion/vgg19-caffe.pth --no-check-certificate -nc https://www.dropbox.com/s/xlbdo688dy4keyk/vgg19-caffe.pth?dl=1')
model = VGGModified(vgg19(pretrained=False), 0.2)
model.load_state_dict(torch.load('vgg_pytorch_modified.pkl')['state_dict'])
return model
else:
assert False
def vgg19(num_classes=1000, pretrained='imagenet'):
"""VGG 19-layer model (configuration "E")
"""
model = models.vgg19(pretrained=False)
if pretrained is not None:
settings = pretrained_settings['vgg19'][pretrained]
model = load_pretrained(model, num_classes, settings)
return model
def contentFunc(self): conv_3_3_layer = 14 cnn = models.vgg19(pretrained=True).features cnn = cnn.cuda() model = nn.Sequential() model = model.cuda() for i,layer in enumerate(list(cnn)): model.add_module(str(i),layer) if i == conv_3_3_layer: break return model
def __init__(self, descriptor_name):
super(Net, self).__init__()
# if descriptor_name == 'vgg16':
# self.select = ['30']
# self.vgg16 = models.vgg16(pretrained=True)
# self.sequence = []
# for name, layer in self.vgg16.features._modules.items():
# self.sequence += [layer]
# for name, layer in self.vgg16.classifier._modules.items():
# self.sequence += [layer]
# break
# self.model = nn.Sequential(*self.sequence)
if descriptor_name == 'vgg16':
self.select = ['30']
self.vgg16 = models.vgg16(pretrained=True)
self.sequence = []
for name, layer in self.vgg16.features._modules.items():
self.sequence += [layer]
for name, layer in self.vgg16.classifier._modules.items():
if name == '6':
break
self.sequence += [layer]
layer = nn.Linear(4096, 10)
# init.xavier_normal(layer.weight.data, gain = 1)
self.sequence += [layer]
self.model = nn.Sequential(*self.sequence)
elif descriptor_name == 'vgg19':
self.select = ['36']
self.vgg19 = models.vgg19(pretrained=True)
self.sequence = []
for name, layer in self.vgg19.features._modules.items():
self.sequence += [layer]
for name, layer in self.vgg19.classifier._modules.items():
self.sequence += [layer]
break
self.model = nn.Sequential(*self.sequence)
elif descriptor_name == 'resnet50':
self.select = ['avgpool']
self.model = models.resnet50(pretrained=True)
self.model.fc = nn.Linear(2048, 10)
elif descriptor_name == 'resnet101':
self.select = ['avgpool']
self.model = models.resnet101(pretrained=True)
elif descriptor_name == 'resnet152':
self.select = ['avgpool']
self.model = models.resnet152(pretrained=True)
self.model.fc = nn.Linear(2048, 10)
def nst_model(content_img, style_img):
vgg = models.vgg19(pretrained=True).features.eval()
normalization = Normalization(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
content_img = content_img.detach()
style_img = style_img.detach()
content_losses = []
style_losses = []
model = nn.Sequential(normalization)
i = 0
for name, layer in vgg._modules.items():
if name in ['0','2','5','10']:
model.add_module('conv_{}'.format(i),layer)
style_target = model(style_img)
style_loss = SL(style_target)
style_losses.append(style_loss)
model.add_module('styleloss_{}'.format(i),style_loss)
i += 1
elif name in ['7']:
model.add_module('conv_{}'.format(i),layer)
content_target = model(content_img)
content_loss = CL(content_target)
content_losses.append(content_loss)
model.add_module('contentloss_{}'.format(i),content_loss)
style_target = model(style_img)
style_loss = SL(style_target)
style_losses.append(style_loss)
model.add_module('styleloss_{}'.format(i),style_loss)
i += 1
elif name in ['1','3','6','8']:
layer = nn.ReLU(inplace=False)
model.add_module('relu_{}'.format(i),layer)
i += 1
elif name in ['4','9']:
model.add_module('maxpool_{}'.format(i),layer)
i += 1
elif name == '11':
break
return model, style_losses, content_losses
def __init__(self, model):
super(feature_net, self).__init__()
if model == 'vgg':
vgg = models.vgg19(pretrained=True)
self.feature == nn.Sequential(*list(vgg.children())[:-1])
self.feature.add_module('global average', nn.AvgPool2d(9))
elif model == 'inceptionv3':
inception = models.inception_v3(pretrained=True)
self.feature = nn.Sequential(*list(inception.children())[:-1])
self.feature._modules.pop('13')
self.feature.add_modules('global average', nn.AvgPool2d(35))
elif model == 'resnet152':
resnet = models.resnet152(pretrained=True)
self.feature = nn.Sequential(*list(resnet.children())[:-1])
def get_vgg19_caffe():
model = vgg19()
model.classifier = nn.Sequential(View(), *model.classifier._modules.values())
vgg = model.features
vgg_classifier = model.classifier
names = ['conv1_1','relu1_1','conv1_2','relu1_2','pool1',
'conv2_1','relu2_1','conv2_2','relu2_2','pool2',
'conv3_1','relu3_1','conv3_2','relu3_2','conv3_3','relu3_3','conv3_4','relu3_4','pool3',
'conv4_1','relu4_1','conv4_2','relu4_2','conv4_3','relu4_3','conv4_4','relu4_4','pool4',
'conv5_1','relu5_1','conv5_2','relu5_2','conv5_3','relu5_3','conv5_4','relu5_4','pool5',
'torch_view','fc6','relu6','drop6','fc7','relu7','drop7','fc8']
model = nn.Sequential()
for n, m in zip(names, list(vgg) + list(vgg_classifier)):
model.add_module(n, m)
model.load_state_dict(torch.load('vgg19-caffe-py3.pth'))
return model
def __init__(self, requires_grad=False):
super(Vgg19, self).__init__()
vgg_pretrained_features = models.vgg19(pretrained=True).features
self.slice1 = torch.nn.Sequential()
self.slice2 = torch.nn.Sequential()
self.slice3 = torch.nn.Sequential()
self.slice4 = torch.nn.Sequential()
self.slice5 = torch.nn.Sequential()
for x in range(2):
self.slice1.add_module(str(x), vgg_pretrained_features[x])
for x in range(2, 7):
self.slice2.add_module(str(x), vgg_pretrained_features[x])
for x in range(7, 12):
self.slice3.add_module(str(x), vgg_pretrained_features[x])
for x in range(12, 21):
self.slice4.add_module(str(x), vgg_pretrained_features[x])
for x in range(21, 30):
self.slice5.add_module(str(x), vgg_pretrained_features[x])
if not requires_grad:
for param in self.parameters():
param.requires_grad = False
def load_model(arch='vgg19', num_labels=102, hidden_units=4096):
# Load a pre-trained model
if arch=='vgg19':
# Load a pre-trained model
model = models.vgg19(pretrained=True)
elif arch=='alexnet':
model = models.alexnet(pretrained=True)
else:
raise ValueError('Unexpected network architecture', arch)
# Freeze its parameters
for param in model.parameters():
param.requires_grad = False
# Features, removing the last layer
features = list(model.classifier.children())[:-1]
# Number of filters in the bottleneck layer
num_filters = model.classifier[len(features)].in_features
# Extend the existing architecture with new layers
features.extend([
nn.Dropout(),
nn.Linear(num_filters, hidden_units),
nn.ReLU(True),
nn.Dropout(),
nn.Linear(hidden_units, hidden_units),
nn.ReLU(True),
nn.Linear(hidden_units, num_labels),
##nn.Softmax(dim=1)
# Please, notice Softmax layer has not been added as per Pytorch answer:
# https://github.com/pytorch/vision/issues/432#issuecomment-368330817
# It is not either included in its transfer learning tutorial:
# https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
])
model.classifier = nn.Sequential(*features)
return model
# print(height_s, width_s)
# print(height_c, width_c)
style_img = utils.image_to_tensor(style_img).unsqueeze(0)
content_img = utils.image_to_tensor(content_img).unsqueeze(0)
style_img = style_img.to(device, torch.float)
content_img = content_img.to(device, torch.float)
# print('content_img size: ', content_img.size())
# utils.show_pic(style_img, 'style image')
# utils.show_pic(content_img, 'content image')
# -------------------------
# Eval() means the parameters of cnn are frozen.
cnn = models.vgg19(pretrained=True).features.to(config.device0).eval()
cnn_normalization_mean = torch.tensor([0.485, 0.456,
0.406]).to(config.device0)
cnn_normalization_std = torch.tensor([0.229, 0.224,
0.225]).to(config.device0)
# Two different initialization ways
input_img = torch.randn(1, 3, height_c, width_c).to(config.device0)
# input_img = content_img.clone()
# print('input_img size: ', input_img.size())
output = run_style_transfer(cnn, cnn_normalization_mean,
cnn_normalization_std, content_img, style_img,
input_img, style_mask_tensor,
content_mask_tensor, L)
print('Style transfer completed')
def reconstruct_image(img_size=224,
test_case=1,
num_epochs=200,
print_iter=10,
save_iter=50):
# Load pre-trained VGG19 model to extract image features
model = models.vgg19(pretrained=True)
if use_gpu:
model = model.cuda(gpu_id)
model.eval()
# Generate a random image which we will optimize
if use_gpu:
recon_img = Variable(
1e-1 * torch.randn(1, 3, img_size, img_size).cuda(gpu_id),
requires_grad=True)
else:
recon_img = Variable(1e-1 * torch.randn(1, 3, img_size, img_size),
requires_grad=True)
# Define optimizer for previously created image
optimizer = optim.SGD([recon_img], lr=1e2, momentum=0.9)
# Decay learning rate by a factor of 0.1 every x epochs
scheduler = lr_scheduler.StepLR(optimizer, step_size=200, gamma=0.1)
# Use deep generator network to get initial image
if use_dgn:
print('Loading deep generator network...')
# Load pre-trained model tokenizer(vocabulary)
dgn = BigGAN.from_pretrained('biggan-deep-256')
# Prepare an input
truncation = 0.4
class_vector = np.zeros((1, 1000), dtype='float32')
noise_vector = truncated_noise_sample(truncation=truncation,
batch_size=1)
# All in tensors
class_vector = torch.from_numpy(class_vector)
noise_vector = torch.from_numpy(noise_vector)
if use_gpu:
class_vector = class_vector.cuda(gpu_id)
noise_vector = noise_vector.cuda(gpu_id)
dgn.cuda(gpu_id)
# Generate image
with torch.no_grad():
output = dgn(noise_vector, class_vector, truncation)
output = output.cpu()
output = nn.functional.interpolate(output,
size=(img_size, img_size),
mode='bilinear',
align_corners=True)
if use_gpu:
recon_img = Variable(output.cuda(gpu_id), requires_grad=True)
else:
recon_img = Variable(output, requires_grad=True)
# Training
for epoch in range(num_epochs):
scheduler.step()
optimizer.zero_grad()
# Get the features from the model of the generated image
output_features = get_features_from_layers(model, recon_img)
# Calculate the losses
euc_loss, alpha_loss, tv_loss, total_loss = get_total_loss(
recon_img, output_features, original_feats)
# Step
total_loss.backward()
optimizer.step()
# Generate image every x iterations
if (epoch + 1) % print_iter == 0:
print('Epoch %d:\tAlpha: %.6f\tTV: %.6f\tEuc: %.6f\tLoss: %.6f' %
(epoch + 1, alpha_loss.data.cpu().numpy(),
tv_loss.data.cpu().numpy(), euc_loss.data.cpu().numpy(),
total_loss.data.cpu().numpy()))
# Save the image every x iterations
if (epoch + 1) % save_iter == 0:
img_sample = torch.squeeze(recon_img.cpu())
im_path = examples_dir + cls + '_' + key + '_eeg.jpg'
save_image(img_sample, im_path, normalize=True)
def style_transfer(style_img,
content_img,
outputpath='./result.png',
num_steps=500,
style_weight=100000,
content_weight=1,
name='test',
loss_dir='losses'):
'''
the main function of neural style transfer
:param style_img: the image with target style you want to transfer to
:param content_img: the original image, to transfer its style while reserve its content
:param outputpath: the path to save image with transferred style
:param num_steps: number of steps to update parameters
:param style_weight: weight of style
:param content_weight: weight of loss
'''
imsize = 512
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([
transforms.Resize(imsize),
transforms.CenterCrop(imsize),
transforms.ToTensor()
])
style_img = image_loader(style_img, transform, device)
content_img = image_loader(content_img, transform, device)
# use the features module of pretrained vgg19
# need the output of the individual convolution layers to measure content and style loss.
cnn = models.vgg19(pretrained=True).features.to(device).eval()
# desired depth layers to compute style/content losses :
content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
input_img = content_img.clone()
output, style_loss, content_loss = run(cnn,
content_layers_default,
style_layers_default,
content_img,
style_img,
input_img,
device,
num_steps=num_steps,
style_weight=style_weight,
content_weight=content_weight)
output = output.detach().cpu().numpy().squeeze(0).transpose([1, 2, 0])
plt.imsave(outputpath, output)
plt.clf()
x = [i * 10 for i in range(len(style_loss))]
plt.plot(x, style_loss, label='style_loss')
plt.plot(x, content_loss, label='content_loss')
plt.xlabel('steps')
plt.ylabel('loss')
plt.legend()
plt.savefig(os.path.join(loss_dir, "loss" + name))
def __init__(self):
super(VGG19, self).__init__()
vgg19 = models.vgg19(pretrained=False)
# Load downloaded parameters
print('Load VGG...', end=' ')
vgg19.load_state_dict(torch.load('./vgg19/vgg19-dcbb9e9d.pth'))
print('Load end.')
features = vgg19.features
self.relu1_1 = torch.nn.Sequential()
self.relu1_2 = torch.nn.Sequential()
self.relu2_1 = torch.nn.Sequential()
self.relu2_2 = torch.nn.Sequential()
self.relu3_1 = torch.nn.Sequential()
self.relu3_2 = torch.nn.Sequential()
self.relu3_3 = torch.nn.Sequential()
self.relu3_4 = torch.nn.Sequential()
self.relu4_1 = torch.nn.Sequential()
self.relu4_2 = torch.nn.Sequential()
self.relu4_3 = torch.nn.Sequential()
self.relu4_4 = torch.nn.Sequential()
self.relu5_1 = torch.nn.Sequential()
self.relu5_2 = torch.nn.Sequential()
self.relu5_3 = torch.nn.Sequential()
self.relu5_4 = torch.nn.Sequential()
for x in range(2):
self.relu1_1.add_module(str(x), features[x])
for x in range(2, 4):
self.relu1_2.add_module(str(x), features[x])
for x in range(4, 7):
self.relu2_1.add_module(str(x), features[x])
for x in range(7, 9):
self.relu2_2.add_module(str(x), features[x])
for x in range(9, 12):
self.relu3_1.add_module(str(x), features[x])
for x in range(12, 14):
self.relu3_2.add_module(str(x), features[x])
for x in range(14, 16):
self.relu3_3.add_module(str(x), features[x])
for x in range(16, 18):
self.relu3_4.add_module(str(x), features[x])
for x in range(18, 21):
self.relu4_1.add_module(str(x), features[x])
for x in range(21, 23):
self.relu4_2.add_module(str(x), features[x])
for x in range(23, 25):
self.relu4_3.add_module(str(x), features[x])
for x in range(25, 27):
self.relu4_4.add_module(str(x), features[x])
for x in range(27, 30):
self.relu5_1.add_module(str(x), features[x])
for x in range(30, 32):
self.relu5_2.add_module(str(x), features[x])
for x in range(32, 34):
self.relu5_3.add_module(str(x), features[x])
for x in range(34, 36):
self.relu5_4.add_module(str(x), features[x])
# don't need the gradients, just want the features
for param in self.parameters():
param.requires_grad = False
def run_train(args):
np.random.seed(args.seed)
torch.manual_seed(args.seed)
print('running training processing...')
style_image = load_image(args.style_image,
mask=False,
size=args.image_style_size,
scale=args.style_scale,
square=True)
style_image = preprocess(style_image)
#save_image('style_image.png',style_image)
cnn = None
if args.loss_model == 'vgg19':
cnn = models.vgg19(pretrained=True).features.to(device).eval()
elif args.loss_model == 'vgg16':
cnn = models.vgg16(pretrained=True).features.to(device).eval()
# get tranform net, content losses, style losses
loss_net, content_losses, style_losses, tv_loss = build_loss_model(
cnn, args, style_image)
#print(loss_net)
#collect space back
cnn = None
del cnn
if args.backend == 'cudnn':
torch.backends.cudnn.enabled = True
#this is to define the inchannels of transferm_net
in_channels = 3
transform = transforms.Compose([
transforms.Resize(args.image_size),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
])
train_dataset = datasets.ImageFolder(args.dataset, transform)
train_loader = DataLoader(train_dataset, batch_size=args.batch_size)
transform_net = TransformerNet(in_channels).to(device)
mse_loss = nn.MSELoss()
optimizer = optim.Adam(transform_net.parameters(), lr=args.learning_rate)
iteration = [0]
while iteration[0] <= args.epochs - 1:
transform_net.train()
agg_content_loss = 0.
agg_style_loss = 0.
count = 0
for batch_id, (x, _) in enumerate(train_loader):
stloss = 0.
ctloss = 0.
n_batch = len(x)
count += n_batch
optimizer.zero_grad()
#stack color_content_masks into x as input
x, x_ori = x.to(device), x.to(device).clone()
x = preprocess(x)
x_ori = preprocess(x_ori)
#save_image('content_x1.png', x[3].unsqueeze(0))
#assert 0 == 1
#forward input to transform_net
y = transform_net(x)
#compute pixel loss
pixloss = 0.
if args.pixel_weight > 0:
pixloss = mse_loss(x, y) * args.pixel_weight
#compute content loss and style loss
for ctl in content_losses:
ctl.mode = 'capture'
for stl in style_losses:
stl.mode = 'None'
loss_net(x_ori)
for ctl in content_losses:
ctl.mode = 'loss'
for stl in style_losses:
stl.mode = 'loss'
loss_net(y)
for ctl in content_losses:
ctloss += mse_loss(ctl.target, ctl.input) * args.content_weight
for stl in style_losses:
local_G = gram_matrix(stl.input)
stloss += mse_loss(local_G, stl.target) * args.style_weight
if tv_loss is not None:
tvloss = tv_loss.loss
else:
tvloss = 0.
loss = ctloss + stloss + pixloss #+ tvloss
loss.backward()
optimizer.step()
agg_content_loss += ctloss.item()
agg_style_loss += stloss.item()
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}, Epoch {}:\t[{}/{}], content: {:.6f}, style: {:.6f}, total: {:.6f}".format(
time.ctime(), iteration[0], count, len(train_dataset),
agg_content_loss / (batch_id + 1),
agg_style_loss / (batch_id + 1),
(agg_content_loss + agg_style_loss) / (batch_id + 1))
print(mesg)
if args.checkpoint_model_dir is not None and (
batch_id + 1) % args.checkpoint_interval == 0:
transform_net.eval().cpu()
ckpt_model_filename = "ckpt_epoch_" + str(
iteration[0] + 1) + "_batch_id_" + str(batch_id +
1) + ".pth"
ckpt_model_path = os.path.join(args.checkpoint_model_dir,
ckpt_model_filename)
torch.save(transform_net.state_dict(), ckpt_model_path)
transform_net.to(device).train()
iteration[0] += 1
#save final model
transform_net.eval().cpu()
save_model_filename = "epoch_" + str(args.epochs) + "_" + str(
time.ctime()).replace(' ', '_') + "_content_" + str(
args.content_weight) + "_style_" + str(
args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
torch.save(transform_net.state_dict(), save_model_path)
print("\n training process is Done!, trained model saved at",
save_model_path)
def main(args):
#with torch.cuda.device(args.gpu):
layers_map = {
'relu4_2': '22',
'relu2_2': '8',
'relu3_2': '13',
'relu1_2': '4'
}
vis = visdom.Visdom(port=args.display_port)
loss_graph = {
"g": [],
"gd": [],
"gf": [],
"gpl": [],
"gpab": [],
"gs": [],
"d": [],
"gdl": [],
"dl": [],
}
# for rgb the change is to feed 3 channels to D instead of just 1. and feed 3 channels to vgg.
# can leave pixel separate between r and gb for now. assume user use the same weights
transforms = get_transforms(args)
if args.color_space == 'rgb':
args.pixel_weight_ab = args.pixel_weight_rgb
args.pixel_weight_l = args.pixel_weight_rgb
rgbify = custom_transforms.toRGB()
train_dataset = ImageFolder('train', args.data_path, transforms)
train_loader = DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True)
val_dataset = ImageFolder('val', args.data_path, transforms)
indices = torch.randperm(len(val_dataset))
val_display_size = args.batch_size
val_display_sampler = SequentialSampler(indices[:val_display_size])
val_loader = DataLoader(dataset=val_dataset,
batch_size=val_display_size,
sampler=val_display_sampler)
# renormalize = transforms.Normalize(mean=[+0.5+0.485, +0.5+0.456, +0.5+0.406], std=[0.229, 0.224, 0.225])
feat_model = models.vgg19(pretrained=True)
netG, netD, netD_local = get_models(args)
criterion_gan, criterion_pixel_l, criterion_pixel_ab, criterion_style, criterion_feat, criterion_texturegan = get_criterions(
args)
real_label = 1
fake_label = 0
optimizerD = optim.Adam(netD.parameters(),
lr=args.learning_rate_D,
betas=(0.5, 0.999))
optimizerG = optim.Adam(netG.parameters(),
lr=args.learning_rate,
betas=(0.5, 0.999))
optimizerD_local = optim.Adam(netD_local.parameters(),
lr=args.learning_rate_D_local,
betas=(0.5, 0.999))
with torch.cuda.device(args.gpu):
netG.cuda()
netD.cuda()
netD_local.cuda()
feat_model.cuda()
criterion_gan.cuda()
criterion_pixel_l.cuda()
criterion_pixel_ab.cuda()
criterion_feat.cuda()
criterion_texturegan.cuda()
input_stack = torch.FloatTensor().cuda()
target_img = torch.FloatTensor().cuda()
target_texture = torch.FloatTensor().cuda()
segment = torch.FloatTensor().cuda()
label = torch.FloatTensor(args.batch_size).cuda()
label_local = torch.FloatTensor(args.batch_size).cuda()
extract_content = FeatureExtractor(feat_model.features,
[layers_map[args.content_layers]])
extract_style = FeatureExtractor(
feat_model.features,
[layers_map[x.strip()] for x in args.style_layers.split(',')])
model = {
"netG": netG,
"netD": netD,
"netD_local": netD_local,
"criterion_gan": criterion_gan,
"criterion_pixel_l": criterion_pixel_l,
"criterion_pixel_ab": criterion_pixel_ab,
"criterion_feat": criterion_feat,
"criterion_style": criterion_style,
"criterion_texturegan": criterion_texturegan,
"real_label": real_label,
"fake_label": fake_label,
"optimizerD": optimizerD,
"optimizerD_local": optimizerD_local,
"optimizerG": optimizerG
}
for epoch in range(args.load_epoch, args.num_epoch):
train(model, train_loader, val_loader, input_stack, target_img,
target_texture, segment, label, label_local, extract_content,
extract_style, loss_graph, vis, epoch, args)
# Importing the Model # ------------------- # # Now we need to import a pre-trained neural network. We will use a 19 # layer VGG network like the one used in the paper. # # PyTorch’s implementation of VGG is a module divided into two child # ``Sequential`` modules: ``features`` (containing convolution and pooling layers), # and ``classifier`` (containing fully connected layers). We will use the # ``features`` module because we need the output of the individual # convolution layers to measure content and style loss. Some layers have # different behavior during training than evaluation, so we must set the # network to evaluation mode using ``.eval()``. # cnn = models.vgg19(pretrained=True).features.to(device).eval() ###################################################################### # Additionally, VGG networks are trained on images with each channel # normalized by mean=[0.485, 0.456, 0.406] and std=[0.229, 0.224, 0.225]. # We will use them to normalize the image before sending it into the network. # cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device) cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device) # create a module to normalize input image so we can easily put it in a # nn.Sequential class Normalization(nn.Module):
def __init__(self):
super(Vgg19, self).__init__()
features = list(vgg19(pretrained=True).features)
self.features = torch.nn.ModuleList(features).eval()
opts = parser.parse_args()
'''
VGG19 test using examples from https://github.com/jacobgil/pytorch-grad-cam for comparison
'''
from torchvision.models import vgg19
if opts.timer:
import time
# this block is mostly copied/adapted from the repo linked above since it's VGG preprocessing stuff
x = cv2.imread(opts.imname, 1) # 1 for color image
x = np.float32(cv2.resize(x, (224, 224))) / 255 # move to range [0,1]
im = preprocess_image(x)
im = torch.cat(opts.batchsize * [im], dim=0)
# create network and set to eval mode
vgg = vgg19(pretrained=True)
# create GradCAM object and generae grad-CAMs
GC = GradCAM(model=vgg, device=opts.device, verbose=opts.verbose)
if opts.classes.__class__.__name__ == 'list':
classes = opts.classes
elif opts.classes.__class__.__name__ == 'int':
classes = [opts.classes]
else:
raise Exception("Invalid list of classes provided!")
# place hooks on ReLU modules inside the vgg.features submodule
if opts.guided:
hookmods = []
for name, module in vgg.features._modules.items():
if module.__class__.__name__ == 'ReLU':
def main(net_type='vgg'):
epoch = 100
batch_size = 64
n_val = 1000
lr = 0.0001
decay = 0.00005
momentum = 0.9
cudnn.benchmark = True
logger.info("Define Network and Loss...")
if net_type == 'vgg':
net = models.vgg19(pretrained=True)
elif net_type == 'resnet':
net = models.resnet50(pretrained=True)
net.to(device)
criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.SGD(net.parameters(),
lr=lr,
momentum=momentum,
weight_decay=decay)
logger.info("Finish")
traindir = '/home/lhy/ILSVRC2012/train'
valdir = '/home/lhy/ILSVRC2012/val'
# train_dataset = NoriInpaintDataset(train_img_nori_list_path, train_img_nori_path, mask_flist_paths_dict)
# train_dataloader = train_dataset.loader(batch_size=batch_size, num_workers=batch_size//4, shuffle=True)
# val_dataset = NoriInpaintDataset(val_img_nori_list_path, val_img_nori_path, mask_flist_paths_dict)
# val_dataloader = val_dataset.loader(batch_size=batch_size, num_workers=batch_size//4, shuffle=False)
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=16,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=batch_size,
shuffle=False,
num_workers=16,
pin_memory=True)
for i in range(epoch):
adjust_learning_rate(opt, epoch, lr)
train(net, train_loader, i, opt, criterion)
# evaluate on validation set
prec1, perc5 = validate(net, val_loader, criterion)
msg = "Epoch {}/{} Val, Perc1:{}, Perc5{}".format(
i, epoch, prec1, perc5)
logger.info(msg)
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': net_type,
'state_dict': net.state_dict(),
'best_prec1': best_prec1,
'optimizer': opt.state_dict(),
},
is_best,
filename=sys.argv[1] + "checkpoint.pth.tar")
def main():
print("SSSRNet3_mask training finetuning on VOC 160*160 patches.")
global opt, model, netContent
opt = parser.parse_args()
print(opt)
gpuid = 0
cuda = True
if cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
if opt.vgg_loss:
print('===> Loading VGG model')
netVGG = models.vgg19()
netVGG.load_state_dict(
model_zoo.load_url(
'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'))
class _content_model(nn.Module):
def __init__(self):
super(_content_model, self).__init__()
self.feature = nn.Sequential(
*list(netVGG.features.children())[:-1])
def forward(self, x):
out = self.feature(x)
return out
netContent = _content_model()
print("===> Building model")
model = Net()
print(get_n_params(model))
model_pretrained = torch.load(
'model/model_DIV2K_noBN_96_epoch_36.pth',
map_location=lambda storage, loc: storage)["model"]
finetune = True
if finetune == True:
index = 0
for (src, dst) in zip(model_pretrained.parameters(),
model.parameters()):
if index > 1:
list(model.parameters())[index].data = src.data
index = index + 1
criterion = nn.MSELoss(size_average=False)
print("===> Setting GPU")
if cuda:
model = model.cuda(gpuid)
model_pretrained = model_pretrained.cuda(gpuid)
criterion = criterion.cuda(gpuid)
if opt.vgg_loss:
netContent = netContent.cuda(gpuid)
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model"].state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training1")
#root_dir = '/tmp4/hang_data/DIV2K/DIV2K_train_320_HDF5'
root_dir = '/tmp4/hang_data/VOCdevkit/VOC2012/VOC_train_label160_HDF5'
files_num = len(os.listdir(root_dir))
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
#save_checkpoint(model, epoch)
print("===> Loading datasets")
x = random.sample(os.listdir(root_dir), files_num)
for index in range(0, files_num):
train_path = os.path.join(root_dir, x[index])
print("===> Training datasets: '{}'".format(train_path))
train_set = DatasetFromHdf5(train_path)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads,
batch_size=opt.batchSize,
shuffle=True)
avgloss = train(training_data_loader, optimizer, model,
model_pretrained, criterion, epoch, gpuid)
if epoch % 2 == 0:
save_checkpoint(model, epoch)
parser.add_argument('--arch', default='vgg19')
parser.add_argument('--epochs', default=5)
parser.add_argument('--hidden_units', default=5000)
parser.add_argument('--learnrate', default=0.003)
parser.add_argument('--gpu', action='store_true')
args = parser.parse_args()
# Load data and create data loaders
train_data, valid_data, test_data = get_data(args.data_dir)
trainloader, validloader, testloader = get_loader(train_data, valid_data,
test_data)
# Load model
if args.arch == 'vgg19':
model = models.vgg19(pretrained=True)
if args.arch == 'vgg13':
model = models.vgg13(pretrained=True)
# Define hyperparameters
learn_rate = float(args.learnrate)
epochs = int(args.epochs)
fc1_out = fc2_in = int(args.hidden_units)
fc2_out = output_in = 1000
# Freeze parameters so we don't backprop through them
for param in model.parameters():
param.requires_grad = False
# Define new classifier to be trained on flower datset with 102 classes
def __init__(self, target_feature):
super(StyleLoss, self).__init__()
self.target = gram_matrix(target_feature).detach()
def forward(self, input):
G = gram_matrix(input)
self.loss = F.mse_loss(G, self.target)
return input
# -----------------------
# Load the neural network
# -----------------------
cnn = models.vgg19(pretrained=True).features.to(device).eval()
cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
# create a module to normalize input image so we can easily put it in a
# nn.Sequential
class Normalization(nn.Module):
def __init__(self, mean, std):
super(Normalization, self).__init__()
# .view the mean and std to make them [C x 1 x 1] so that they can
# directly work with image Tensor of shape [B x C x H x W].
# B is batch size. C is number of channels. H is height and W is width.
self.mean = torch.tensor(mean).view(-1, 1, 1)
self.std = torch.tensor(std).view(-1, 1, 1)
def __init__(self):
super(Model, self).__init__()
self.backbone = models.vgg19()
num_features = self.backbone.classifier[6].in_features
self.backbone.classifier[6] = nn.Linear(num_features, len(classes))
def initialModel(modelName, numClasses, featureExtract, usePretrained):
if modelName == "vgg19":
modelUse = models.vgg19(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0'] = nn.Conv2d(1,
64,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1))
modelUse.classifier._modules['6'] = nn.Linear(in_features=4096,
out_features=numClasses,
bias=True)
elif modelName == "shufflenet_v2_x2_0":
modelUse = models.shufflenet_v2_x2_0(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(2048, numClasses, bias=True))
else:
modelUse.fc = nn.Linear(2048, numClasses, bias=True)
elif modelName == "shufflenet_v2_x0_5":
modelUse = models.shufflenet_v2_x0_5(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(
nn.Dropout(dropNum),
nn.Linear(1024, numClasses, bias=True),
)
else:
modelUse.fc = nn.Linear(1024, numClasses, bias=True)
elif modelName == "shufflenet_v2_x1_5":
modelUse = models.shufflenet_v2_x1_5(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(
nn.Dropout(dropNum),
nn.Linear(1024, numClasses, bias=True),
)
else:
modelUse.fc = nn.Linear(1024, numClasses, bias=True)
elif modelName == "shufflenet_v2_x1_0":
modelUse = models.shufflenet_v2_x1_0(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.conv1._modules['0'] = nn.Conv2d(1,
24,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(1024, numClasses, bias=True))
else:
modelUse.fc = nn.Linear(1024, numClasses, bias=True)
elif modelName == "squeezenet1_1":
modelUse = models.squeezenet1_1(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0'] = nn.Conv2d(1,
64,
kernel_size=(3, 3),
stride=(2, 2))
modelUse.classifier = nn.Sequential(
nn.Dropout(p=dropNum, inplace=False),
nn.Conv2d(512, numClasses, kernel_size=(1, 1), stride=(1, 1)),
nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(output_size=(1, 1)))
elif modelName == "densenet121":
modelUse = models.densenet121(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['conv0'] = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.classifier = nn.Sequential(
nn.Dropout(dropNum), nn.Linear(1024, numClasses, bias=True))
else:
modelUse.classifier = nn.Linear(1024, numClasses, bias=True)
elif modelName == "vgg16":
modelUse = models.vgg16(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0'] = nn.Conv2d(1,
64,
kernel_size=(3, 3),
stride=(1, 1),
padding=(1, 1))
modelUse.classifier._modules['6'] = nn.Linear(in_features=4096,
out_features=numClasses,
bias=True)
elif modelName == 'mobilenet_v2':
modelUse = models.mobilenet_v2(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
if mono:
modelUse.features._modules['0']._modules['0'] = nn.Conv2d(
1,
32,
kernel_size=(3, 3),
stride=(2, 2),
padding=(1, 1),
bias=False)
modelUse.classifier = nn.Sequential(
nn.Dropout(p=dropNum, inplace=False),
nn.Linear(in_features=1280, out_features=numClasses, bias=True))
elif modelName == "resnet18":
modelUse = models.resnet18(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
numFeatures = modelUse.fc.in_features
if mono:
modelUse.conv1 = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(numFeatures, numClasses))
else:
modelUse.fc = nn.Linear(numFeatures, numClasses)
elif modelName == "resnet34":
modelUse = models.resnet34(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
numFeatures = modelUse.fc.in_features
if mono:
modelUse.conv1 = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(numFeatures, numClasses))
else:
modelUse.fc = nn.Linear(numFeatures, numClasses)
elif modelName == "resnet50":
modelUse = models.resnet50(pretrained=usePretrained)
setParameterRequiresGrad(modelUse, featureExtract)
numFeatures = modelUse.fc.in_features
if mono:
modelUse.conv1 = nn.Conv2d(1,
64,
kernel_size=(7, 7),
stride=(2, 2),
padding=(3, 3),
bias=False)
if dropout:
modelUse.fc = nn.Sequential(nn.Dropout(dropNum),
nn.Linear(numFeatures, numClasses))
else:
modelUse.fc = nn.Linear(numFeatures, numClasses)
else:
print("model not implemented")
return None, None
return modelUse
def train_model(args, device, cfg=cfg):
# Load VGG19 features
vgg = vgg19(pretrained=True).features
vgg = vgg.to(device)
# We don't want to train VGG
for param in vgg.parameters():
param.requires_grad_(False)
# Load style net
style_net = HRNet()
style_net = style_net.to(device)
# Load images
content_img = load_image(os.path.join(args.img_root, args.content_img), size=args.content_size)
content_img = content_img.to(device)
style_img = load_image(os.path.join(args.img_root, args.style_img))
style_img = style_img.to(device)
# Get features from VGG
content_features = get_features(content_img, vgg)
style_features = get_features(style_img, vgg)
target = content_img.clone().requires_grad_(True).to(device)
style_gram_matrixs = {layer: get_gram_matrix(style_features[layer]) for layer in style_features}
optim = torch.optim.Adam(style_net.parameters(), lr=cfg['lr'])
scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=cfg['step_size'], gamma=cfg['gamma'])
content_loss_epoch = []
style_loss_epoch = []
total_loss_epoch = []
i = 0
output_image = content_img
content_weight = args.content_weight
style_weight = args.style_weight
# Start training
for epoch in tqdm(range(cfg['steps']+1)):
scheduler.step()
target = style_net(content_img).to(device)
target.requires_grad_(True)
target_features = get_features(target, vgg)
content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2']) ** 2)
style_loss = 0
for layer in style_weights:
target_feature = target_features[layer]
target_gram_matrix = get_gram_matrix(target_feature)
style_gram_matrix = style_gram_matrixs[layer]
layer_style_loss = style_weights[layer] * torch.mean((target_gram_matrix - style_gram_matrix) ** 2)
b, c, h, w = target_feature.shape
style_loss += layer_style_loss / (c*h*w)
total_loss = content_weight * content_loss + style_weight * style_loss
total_loss_epoch.append(total_loss.item())
style_loss_epoch.append(style_weight * style_loss)
content_loss_epoch.append(content_weight * content_loss.item())
optim.zero_grad()
total_loss.backward()
optim.step()
if epoch % cfg['show_every'] == 0:
print("After %d criterions:" % epoch)
print('Total loss: ', total_loss.item())
print('Content loss: ', content_loss.item())
print('Style loss: ', style_loss.item())
plt.imshow(im_convert(target))
plt.show()
# plt.savefig(f'{args.save_dir}fig{i}.png')
# i += 1
output_image = target
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
root_dir = 'dataset'
# train_dir = root_dir + '/train'
# valid_dir = root_dir + '/valid'
batch_size = 32
# Define your transforms for the training and testing sets
data_transforms = {
'train':
transforms.Compose([
transforms.RandomRotation(30),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'valid':
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
}
# Load the datasets with ImageFolder
image_datasets = {
x: datasets.ImageFolder(os.path.join(root_dir, x), data_transforms[x])
for x in ['train', 'valid']
}
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True,
num_workers=4)
for x in ['train', 'valid']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}
class_names = image_datasets['train'].classes
num_class = len(class_names)
print(dataset_sizes)
# # Label mapping
# with open('categories.json', 'r') as f:
# categories: object = json.load(f)
# # Run this to test the data loader
# images, labels = next(iter(dataloaders['train']))
# rand_idx = np.random.randint(len(images))
# # Print(rand_idx)
# print("label: {}, class: {}, name: {}".format(labels[rand_idx].item(),
# class_names[labels[rand_idx].item()],
# categories[class_names[labels[rand_idx].item()]]))
def build_classifier(input_features, hidden_layers, output_features):
classifier = nn.conv2d()
if hidden_layers is None:
classifier.add_module('fc0',
nn.conv2d(input_features, output_features))
else:
layer_sizes = zip(hidden_layers[:-1], hidden_layers[1:])
classifier.add_module('fc0',
nn.conv2d(input_features, hidden_layers[0]))
classifier.add_module('relu0', nn.ReLU())
classifier.add_module('drop0', nn.Dropout(.6))
classifier.add_module('relu1', nn.ReLU())
classifier.add_module('drop1', nn.Dropout(.5))
for i, (h1, h2) in enumerate(layer_sizes):
classifier.add_module('fc' + str(i + 1), nn.Linear(h1, h2))
classifier.add_module('relu' + str(i + 1), nn.ReLU())
classifier.add_module('drop' + str(i + 1), nn.Dropout(.5))
classifier.add_module(
'output', nn.Linear(hidden_layers[-1], output_features))
return classifier
def train_model(model, criterion, optimizer, sched, num_epochs):
since = time.time()
best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch + 1, num_epochs))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train', 'valid']:
if phase == 'train':
print("in training mode:")
model.train() # Set model to training mode
else:
print("in evaluation mode:")
model.eval() # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
print('iterating over data')
# zero the parameter gradients
optimizer.zero_grad()
# forward
# track history if only in train
print('forward propagation')
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
print('backward propagation')
sched.step()
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / dataset_sizes[phase]
epoch_acc = running_corrects.double() / dataset_sizes[phase]
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
# deep copy the model
if phase == 'valid' and epoch_acc > best_acc:
best_acc = epoch_acc
best_model_wts = copy.deepcopy(model.state_dict())
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print('Best val Acc: {:4f}'.format(best_acc))
# load best model weights
model.load_state_dict(best_model_wts)
return model
# Create classifier
model = models.vgg19(pretrained=True)
for parameter in model.parameters():
parameter.requires_grad = False
num_in_features = 25088
num_hidden_layers = None
num_out_features = 9
new_classifier = build_classifier(num_in_features, num_hidden_layers,
num_out_features)
# model.classifier = new_classifier
# train model
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=0.0001)
scheduler = lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)
epochs = 1
model.to(device)
model = train_model(model, criterion, optimizer, scheduler, epochs)
# Evaluation
model.eval()
accuracy = 0
for inputs, labels in dataloaders['valid']:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
# Class with the highest probability is our predicted class
equality = (labels.data == outputs.max(1)[1])
# Accuracy is number of correct predictions divided by all predictions
accuracy += equality.type_as(torch.FloatTensor()).mean()
print("Test accuracy: {:.3f}".format(accuracy / len(dataloaders['valid'])))
model.class_to_idx = image_datasets['train'].class_to_idx
checkpoint = {
'input_size': dataset_sizes['train'],
'output_size': num_class,
'epochs': epochs,
'batch_size': batch_size,
'model': models.vgg19(pretrained=True),
'classifier': new_classifier,
'scheduler': scheduler,
'optimizer': optimizer.state_dict(),
'state_dict': model.state_dict(),
'class_to_idx': model.class_to_idx
}
torch.save(checkpoint, 'trained_model.pth')
def train_save_network(data_dir, save_dir, arch, learning_rate, hidden_units,
epochs, gpu):
data_dir = data_dir
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
test_dir = data_dir + '/test'
train_transforms = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
train_data = datasets.ImageFolder(train_dir, transform=train_transforms)
trainloader = torch.utils.data.DataLoader(train_data,
batch_size=64,
shuffle=True)
valid_transforms = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
valid_data = datasets.ImageFolder(valid_dir, transform=valid_transforms)
validloader = torch.utils.data.DataLoader(valid_data, batch_size=64)
if arch == 'vgg11':
model = models.vgg11(pretrained=True)
if arch == 'vgg13':
model = models.vgg13(pretrained=True)
if arch == 'vgg16':
model = models.vgg16(pretrained=True)
if arch == 'vgg19':
model = models.vgg19(pretrained=True)
for param in model.parameters():
param.requires_grad = False
classifier = nn.Sequential(
OrderedDict([('fc1', nn.Linear(25088, hidden_units)),
('relu', nn.ReLU()), ('fc2', nn.Linear(hidden_units,
102)),
('output', nn.LogSoftmax(dim=1))]))
model.classifier = classifier
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)
if gpu == True:
device = 'cuda'
else:
device = 'cpu'
epochs = epochs
print_every = 40
steps = 0
model.to(device)
for e in range(epochs):
running_loss = 0
for ii, (inputs, labels) in enumerate(trainloader):
steps += 1
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model.forward(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if steps % print_every == 0:
model.eval()
with torch.no_grad():
valid_loss = 0
accuracy = 0
for images, labels in validloader:
images, labels = images.to(device), labels.to(device)
output = model.forward(images)
valid_loss += criterion(output, labels).item()
ps = torch.exp(output)
equality = (labels.data == ps.max(dim=1)[1])
accuracy += equality.type(torch.FloatTensor).mean()
print(
"Epoch: {}/{}.. ".format(e + 1, epochs),
"Training Loss: {:.3f}.. ".format(running_loss /
print_every),
"Test Loss: {:.3f}.. ".format(valid_loss /
len(validloader)),
"Test Accuracy: {:.3f}".format(accuracy /
len(validloader)))
running_loss = 0
if save_dir == 'none':
pass
else:
model.class_to_idx = train_data.class_to_idx
checkpoint = {
'input_size': 25088,
'output_size': 102,
'hidden_layers': [hidden_units],
'state_dict': model.state_dict(),
'class_to_idx': model.class_to_idx,
'optimizer_state_dict': optimizer.state_dict
}
torch.save(checkpoint, save_dir)
def __init__(self):
"""Select conv1_1 ~ conv5_1 activation maps."""
super(VGGNet, self).__init__()
self.select = ['0', '5', '10', '19', '28']
self.vgg = models.vgg19(pretrained=True).features
def main():
global opt, model, netContent
opt = parser.parse_args()
print(opt)
cuda = opt.cuda
if cuda:
print("=> use gpu id: '{}'".format(opt.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = opt.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
opt.seed = random.randint(1, 10000)
print("Random Seed: ", opt.seed)
torch.manual_seed(opt.seed)
if cuda:
torch.cuda.manual_seed(opt.seed)
cudnn.benchmark = True
print("===> Loading datasets")
training_data_loaders = []
random_order = [i for i in range(4)]
for i in range(4):
filename = "../train/DIV2K_train_320_HDF5/DIV2K_x4_Part" + str(
i + 1) + ".h5"
train_set = DatasetFromHdf5(filename)
training_data_loader = DataLoader(dataset=train_set,
num_workers=opt.threads, \
batch_size=opt.batchSize,
shuffle=True)
training_data_loaders.append(training_data_loader)
if opt.vgg_loss:
print('===> Loading VGG model')
netVGG = models.vgg19()
netVGG.load_state_dict(
model_zoo.load_url(
'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'))
class _content_model(nn.Module):
def __init__(self):
super(_content_model, self).__init__()
self.feature = nn.Sequential(
*list(netVGG.features.children())[:-1])
def forward(self, x):
out = self.feature(x)
return out
netContent = _content_model()
print("===> Building model")
model = _NetG()
criterion = nn.MSELoss(size_average=False)
print("===> Setting GPU")
if cuda:
model = model.cuda()
criterion = criterion.cuda()
if opt.vgg_loss:
netContent = netContent.cuda()
# optionally resume from a checkpoint
if opt.resume:
if os.path.isfile(opt.resume):
print("=> loading checkpoint '{}'".format(opt.resume))
checkpoint = torch.load(opt.resume)
opt.start_epoch = checkpoint["epoch"] + 1
model.load_state_dict(checkpoint["model"].state_dict())
else:
print("=> no checkpoint found at '{}'".format(opt.resume))
# optionally copy weights from a checkpoint
if opt.pretrained:
if os.path.isfile(opt.pretrained):
print("=> loading model '{}'".format(opt.pretrained))
weights = torch.load(opt.pretrained)
model.load_state_dict(weights['model'].state_dict())
else:
print("=> no model found at '{}'".format(opt.pretrained))
print("===> Setting Optimizer")
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
print("===> Training")
for epoch in range(opt.start_epoch, opt.nEpochs + 1):
shuffle(random_order)
for idx in random_order:
train(training_data_loaders[idx], optimizer, model, criterion,
epoch)
save_checkpoint(model, epoch)
import torch
import torch.nn as nn
import torchvision.models as models
import loss
vgg = models.vgg19(pretrained=True).features
if torch.cuda.is_available():
vgg = vgg.cuda()
content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
def get_style_model_and_loss(style_img,
content_img,
cnn=vgg,
style_weight=1000,
content_weight=1,
content_layers=content_layers_default,
style_layers=style_layers_default):
content_loss_list = []
style_loss_list = []
model = nn.Sequential()
if torch.cuda.is_available():
model = model.cuda()
gram = loss.Gram()
if torch.cuda.is_available():
gram = gram.cuda()
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--use-cuda', action='store_true', default=False)
parser.add_argument('--image-path',
type=str,
default='./examples/000004.jpg')
args = parser.parse_args()
args.use_cuda = args.use_cuda and torch.cuda.is_available()
return args
if __name__ == '__main__':
args = get_args()
grad_cam = GradCam(model=models.vgg19(pretrained=True), \
target_layer_names=["35"], use_cuda=args.use_cuda)
img_name = args.image_path.split('/')[-1].split('.')[0]
img = cv2.imread(args.image_path, 1)
img = np.float32(cv2.resize(img, (224, 224))) / 255
input = preprocess_image(img)
# Find tok and plot
noutput = grad_cam.forward(input).cpu().data.numpy()[0]
topk_idxs = noutput.argsort()[-5:][::-1]
print(noutput[topk_idxs])
# If None, returns the map for the highest scoring category.
# Otherwise, targets the requested index.
# target_index = None
for topl, target_index in enumerate(topk_idxs):
print(f'{topl}:{imgnet_classes[target_index]}')
def __init__(self):
super(MeanActivations, self).__init__()
self.vgg = models.vgg19(pretrained=True).features
def __init__(self):
self.cnn = models.vgg19(pretrained=True).features.to(device).eval()
self.num_steps = 70
self.image_size = 100
self.progress_lambda = None
self.should_terminate_lambda = None
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
#----------------------------------------------------------------------------------------------
import argparse
import os
from six import text_type as _text_type
from mmdnn.conversion.examples.imagenet_test import TestKit
import torch
import torchvision.models as models
NETWORKS_MAP = {
'inception_v3' : lambda : models.inception_v3(pretrained=True),
'vgg16' : lambda : models.vgg16(pretrained=True),
'vgg19' : lambda : models.vgg19(pretrained=True),
'resnet152' : lambda : models.resnet152(pretrained=True),
'densenet' : lambda : models.densenet201(pretrained=True),
'squeezenet' : lambda : models.squeezenet1_1(pretrained=True)
}
def _main():
parser = argparse.ArgumentParser()
parser.add_argument('-n', '--network',
type=_text_type, help='Model Type', required=True,
choices=NETWORKS_MAP.keys())
parser.add_argument('-i', '--image', type=_text_type, help='Test Image Path')
def main():
# parse options
# parser = TrainShapeMatchingOptions()
# opts = parser.parse()
# create model
print('--- create model ---')
netShapeM = ShapeMatchingGAN(opts.GS_nlayers, opts.DS_nlayers, opts.GS_nf,
opts.DS_nf, opts.GT_nlayers, opts.DT_nlayers,
opts.GT_nf, opts.DT_nf, opts.gpu)
if opts.gpu:
netShapeM.cuda()
netShapeM.init_networks(weights_init)
netShapeM.train()
if opts.style_loss:
netShapeM.G_S.load_state_dict(torch.load(opts.load_GS_name))
netShapeM.G_S.eval()
VGGNet = models.vgg19(pretrained=True).features
VGGfeatures = VGGFeature(VGGNet, opts.gpu)
for param in VGGfeatures.parameters():
param.requires_grad = False
if opts.gpu:
VGGfeatures.cuda()
style_targets = get_GRAM(opts.style_name, VGGfeatures, opts.batchsize,
opts.gpu)
print('--- training ---')
# load image pair
_, X, Y, Noise = load_style_image_pair(opts.style_name, gpu=opts.gpu)
Y = to_var(Y) if opts.gpu else Y
X = to_var(X) if opts.gpu else X
Noise = to_var(Noise) if opts.gpu else Noise
for epoch in range(opts.texture_step1_epochs):
for i in range(opts.Ttraining_num // opts.batchsize):
x, y = cropping_training_batches(X, Y, Noise, opts.batchsize,
opts.Tanglejitter,
opts.subimg_size,
opts.subimg_size)
losses = netShapeM.texture_one_pass(x, y)
print('Step1, Epoch [%02d/%02d][%03d/%03d]' %
(epoch + 1, opts.texture_step1_epochs, i + 1,
opts.Ttraining_num // opts.batchsize),
end=': ')
print('LDadv: %+.3f, LGadv: %+.3f, Lrec: %+.3f, Lsty: %+.3f' %
(losses[0], losses[1], losses[2], losses[3]))
if opts.style_loss:
fnames, _ = custom_load_train_batchfnames(opts.text_path,
opts.augment_text_path,
opts.batchsize,
opts.text_datasize,
opts.augment_text_datasize,
trainnum=opts.Ttraining_num)
for epoch in range(opts.texture_step2_epochs):
itr = 0
for fname in fnames:
itr += 1
t = prepare_text_batch(fname, anglejitter=False)
x, y = cropping_training_batches(X, Y, Noise, opts.batchsize,
opts.Tanglejitter,
opts.subimg_size,
opts.subimg_size)
t = to_var(t) if opts.gpu else t
losses = netShapeM.texture_one_pass(x, y, t, 0, VGGfeatures,
style_targets)
print('Step2, Epoch [%02d/%02d][%03d/%03d]' %
(epoch + 1, opts.texture_step2_epochs, itr, len(fnames)),
end=': ')
print('LDadv: %+.3f, LGadv: %+.3f, Lrec: %+.3f, Lsty: %+.3f' %
(losses[0], losses[1], losses[2], losses[3]))
print('--- save ---')
# directory
netShapeM.save_texture_model(opts.save_path, opts.save_name)
def __init__(self):
"""Select conv1_1 ~ conv5_1 activation maps."""
super(VGGNet, self).__init__()
self.select = ['0', '5', '10', '19', '28']
self.vgg = models.vgg19(pretrained=True).features
def __init__(self):
super(VGG19, self).__init__()
features = models.vgg19(pretrained=True).features
self.relu1_1 = torch.nn.Sequential()
self.relu1_2 = torch.nn.Sequential()
self.relu2_1 = torch.nn.Sequential()
self.relu2_2 = torch.nn.Sequential()
self.relu3_1 = torch.nn.Sequential()
self.relu3_2 = torch.nn.Sequential()
self.relu3_3 = torch.nn.Sequential()
self.relu3_4 = torch.nn.Sequential()
self.relu4_1 = torch.nn.Sequential()
self.relu4_2 = torch.nn.Sequential()
self.relu4_3 = torch.nn.Sequential()
self.relu4_4 = torch.nn.Sequential()
self.relu5_1 = torch.nn.Sequential()
self.relu5_2 = torch.nn.Sequential()
self.relu5_3 = torch.nn.Sequential()
self.relu5_4 = torch.nn.Sequential()
for x in range(2):
self.relu1_1.add_module(str(x), features[x])
for x in range(2, 4):
self.relu1_2.add_module(str(x), features[x])
for x in range(4, 7):
self.relu2_1.add_module(str(x), features[x])
for x in range(7, 9):
self.relu2_2.add_module(str(x), features[x])
for x in range(9, 12):
self.relu3_1.add_module(str(x), features[x])
for x in range(12, 14):
self.relu3_2.add_module(str(x), features[x])
for x in range(14, 16):
self.relu3_3.add_module(str(x), features[x])
for x in range(16, 18):
self.relu3_4.add_module(str(x), features[x])
for x in range(18, 21):
self.relu4_1.add_module(str(x), features[x])
for x in range(21, 23):
self.relu4_2.add_module(str(x), features[x])
for x in range(23, 25):
self.relu4_3.add_module(str(x), features[x])
for x in range(25, 27):
self.relu4_4.add_module(str(x), features[x])
for x in range(27, 30):
self.relu5_1.add_module(str(x), features[x])
for x in range(30, 32):
self.relu5_2.add_module(str(x), features[x])
for x in range(32, 34):
self.relu5_3.add_module(str(x), features[x])
for x in range(34, 36):
self.relu5_4.add_module(str(x), features[x])
# don't need the gradients, just want the features
for param in self.parameters():
param.requires_grad = False
import sys
from coco_dataset import *
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision.models as models
model = models.vgg19(pretrained=True)
new_classifier = nn.Sequential(*list(model.classifier.children())[:-1])
model.classifier = new_classifier
model.cuda()
mytrans = transforms.Compose([Rescale((224,224)),ToVGGTensor()])
train_dir = '/home/datasets/coco/raw/train2014/'
val_dir = '/home/datasets/coco/raw/val2014'
#dataset = COCOImageDataset(train_dir,val_dir,transform=default_transform(224))
dataset = COCOImageCropDataset(train_dir,val_dir,transform=croped_transform())
dataloader = data.DataLoader(dataset, batch_size=64, shuffle=False, num_workers=4)
for i, batch in enumerate(dataloader):
fns,imgs = batch
imgs = imgs.cuda()
fns = list(fns)
feats = []
for j,fn in enumerate(fns):
feat = model.forward(Variable(imgs[j,:]))
feat = torch.mean(feat.data,dim=0)
feats.append(feat)
parser.add_argument('--style_weight', '-s_w', type=int, default=500, help='The weight of style loss')
parser.add_argument('--initialize_noise', '-i_n', action='store_true', help='Initialize with white noise? elif initialize with content image')
parser.add_argument('--cuda', action='store_true',default=True, help='use cuda?')
args = parser.parse_args()
use_cuda = torch.cuda.is_available() and args.cuda
dtype = torch.cuda.FloatTensor if use_cuda else torch.FloatTensor
# desired size of the output image
imsize = 512 if use_cuda else 128 # use small size if no gpu
#assert style_img.size() == content_img.size(), \
# "we need to import style and content images of the same size"
cnn = models.vgg19(pretrained=True).features
# move it to the GPU if possible:
if use_cuda:
cnn = cnn.cuda()
# desired depth layers to compute style/content losses :
content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
def get_style_model_and_losses(cnn, style_img, content_img, style_weight=1000, content_weight=1, content_layers=content_layers_default, style_layers=style_layers_default):
cnn = copy.deepcopy(cnn)
# just in order to have an iterable access to or list of content/syle
# losses
def train_model(dat_dir, arch, hid_unit, lr, epoch, device, save_name):
#check whether inputs are valid
if os.path.isdir(dat_dir) == False:
print('Please enter a valid path to the pictures')
return
if (hid_unit < 102) | (hid_unit > 25088):
print(
'Please enter a valid number of Hidden Units: 102 < Hid_Units < 25088'
)
return
if lr <= 0:
print('Learning Rate must be >0')
return
if epoch < 0:
print('Epochs number must be >0')
return
if (device != 'cpu') & (device != 'cuda'):
print('Please choose "cpu"/"cuda" for --gpu')
return
# Imports
import numpy as np
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models
data_dir = dat_dir
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
test_dir = data_dir + '/test'
# TODO: Define your transforms for the training, validation, and testing sets
train_transforms = transforms.Compose([
transforms.RandomRotation(30),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
validation_transforms = transforms.Compose([
transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_transforms = transforms.Compose([
transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# TODO: Load the datasets with ImageFolder
train_data = datasets.ImageFolder(train_dir, transform=train_transforms)
validation_data = datasets.ImageFolder(valid_dir,
transform=train_transforms)
test_data = datasets.ImageFolder(test_dir, transform=test_transforms)
# TODO: Using the image datasets and the trainforms, define the dataloaders
trainloader = torch.utils.data.DataLoader(train_data,
batch_size=64,
shuffle=True)
validationloader = torch.utils.data.DataLoader(validation_data,
batch_size=64,
shuffle=False)
testloader = torch.utils.data.DataLoader(test_data,
batch_size=64,
shuffle=False)
if arch == 'vgg13':
model = models.vgg13(pretrained=True)
elif arch == 'vgg16':
model = models.vgg16(pretrained=True)
elif arch == 'vgg19':
model = models.vgg19(pretrained=True)
else:
print('Please choose VGG13/VGG16/VGG19 model')
return
# Freeze model parameters and weights
for param in model.parameters():
param.requires_grad = False
# Define 3 output layers
from collections import OrderedDict
classifier = nn.Sequential(
OrderedDict([('fc1', nn.Linear(25088, hid_unit)), ('relu1', nn.ReLU()),
('dropout1', nn.Dropout(p=0.2)),
('fc4', nn.Linear(hid_unit, 102)),
('output', nn.LogSoftmax(dim=1))]))
# Assign new output layers to a model
model.classifier = classifier
#Define Loss Criterion and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.classifier.parameters(), lr=lr)
# Send Model to Device
model.to(device)
# Loop over Epochs
for e in range(epoch):
# Null the Stat variables
running_loss = 0
test_loss = 0
accuracy = 0
# Loop over Training Batches
for images, labels in trainloader:
# Switch on Train Mode to use gradients
model.train()
#Send images and labels to GPU/CPU
images, labels = images.to(device), labels.to(device)
# NN Routine
optimizer.zero_grad()
output = model.forward(images)
loss = criterion(output, labels)
loss.backward()
optimizer.step()
# Calculate Stat
running_loss += loss.item()
# Validation Step
else:
#Switch off the Gradient calculations (to increase speed)
with torch.no_grad():
# Evaluation Regime (to increase speed)
model.eval()
#Loop over test Batches
for images1, labels1 in validationloader:
#Send images and labels to GPU / CPU
images1, labels1 = images1.to(device), labels1.to(device)
#Forward pass through the model and get the probabilities
ps = torch.exp(model(images1))
#Get classes and probabilities from ts
top_p, top_class = ps.topk(1, dim=1)
#Calculate Loss
loss = criterion(model(images1), labels1)
# Statistics
test_loss += loss.item()
equals = top_class == labels1.view(*top_class.shape)
accuracy += torch.mean(equals.type(torch.FloatTensor))
print(f'Epoch No: {e+1}')
print(
f'Accuracy Validation: {accuracy.item()/len(validationloader)*100}%'
)
print(f'Training Loss: {running_loss/len(trainloader)}')
print(f'Validation Loss: {test_loss/len(validationloader)}')
#Save the model
model.class_to_idx = train_data.class_to_idx
torch.save(
{
'architecture': arch,
'classifier': model.classifier,
'state_dict': model.state_dict(),
'class_to_idx': model.class_to_idx
}, save_name)
def test_vgg19(self):
process_model(models.vgg19(self.pretrained), self.image,
_C_tests.forward_vgg19, "VGG19")
def __init__(self):
super(VGGNet, self).__init__()
self.vgg = models.vgg19(pretrained=True).features
opt = parser.parse_args()
EPOCH = opt.epochs
DIM = opt.input_dimension
BATCH_SIZE = opt.batch_size
LR = opt.learning_rate
cuda_available = torch.cuda.is_available()
device = torch.device("cuda" if cuda_available else "cpu")
print('Using device:%s' % (device))
content_layers = ['conv_5'] # HAO['conv_4']
style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
vgg = vgg19(pretrained=True).to(device)
vgg.eval()
for p in vgg.parameters():
p.requires_grad = False
def checkpoint_path(epoch, batch):
if batch > 0:
return '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/worker_e_%d_b_%d' % (
epoch, batch
), '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/judge_e_%d_b_%d' % (
epoch, batch)
else:
return '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/worker_e_%d' % (
epoch
), '/media/hao/bigfatdrive/FeaturePolice/Checkpoints/judge_e_%d' % (
def preprocess(self, img_path):
temp_img = Image.open(img_path)
return self.transform(temp_img)
def __getitem__(self, index):
sr_img = self.preprocess(
os.path.join(self.sr_path, self.sr_names[index]))
gt_img = self.preprocess(
os.path.join(self.gt_path, self.gt_names[index]))
return sr_img.to(device), gt_img.to(device)
def __len__(self):
return len(self.sr_names)
vgg19 = models.vgg19(pretrained=True).to(device)
vgg19.eval()
vgg_data = VggDataset(sr_path='../results/RRDB_ESRGAN_x4/div-2k-test',
gt_path='../data_samples/div_train/valid_HR')
vgg_loader = DataLoader(vgg_data, batch_size=2)
for sr, gt in vgg_loader:
sr_vgg = vgg19(sr)
gt_vgg = vgg19(gt)
vgg_loss = ((gt_vgg - sr_vgg)**2).mean()
print(vgg_loss)
def __init__(self, config=GlobalConfig()):
super(HighQualityNet, self).__init__()
self.config = config
""" Settings """
if self.config.architecture == 'AlexNet':
self.classification_net = models.alexnet(pretrained=True).cuda()
print(self.classification_net)
elif self.config.architecture == 'ResNet':
self.classification_net = models.resnet50(pretrained=True).cuda()
print(self.classification_net)
elif self.config.architecture == 'VGG':
self.classification_net = models.vgg19(pretrained=True).cuda()
print(self.classification_net)
elif self.config.architecture == 'DenseNet':
self.classification_net = models.densenet121(
pretrained=True).cuda()
print(self.classification_net)
elif self.config.architecture == 'ResNet':
self.classification_net = models.resnet152(pretrained=True).cuda()
print(self.classification_net)
elif self.config.architecture == 'GoogleNet':
self.classification_net = models.googlenet(pretrained=True).cuda()
print(self.classification_net)
else:
self.classification_net = models.mobilenet_v2(
pretrained=True).cuda()
print(self.classification_net)
if torch.cuda.device_count() > 1:
self.classification_net = torch.nn.DataParallel(
self.classification_net)
self.criterion = nn.CrossEntropyLoss().cuda()
self.encoder = Prep_pureUnet(config=config).cuda()
if torch.cuda.device_count() > 1:
self.encoder = torch.nn.DataParallel(self.encoder)
print(self.encoder)
self.optimizer = torch.optim.Adam(self.encoder.parameters())
""" Noise Layers """
self.noise_layers = [Identity()]
# self.cropout_layer = Cropout(config).cuda()
self.jpeg_layer_80 = DiffJPEG(256,
256,
quality=80,
differentiable=True).cuda()
self.jpeg_layer_90 = DiffJPEG(256,
256,
quality=90,
differentiable=True).cuda()
self.jpeg_layer_70 = DiffJPEG(256,
256,
quality=70,
differentiable=True).cuda()
self.jpeg_layer_60 = DiffJPEG(256,
256,
quality=60,
differentiable=True).cuda()
self.jpeg_layer_50 = DiffJPEG(256,
256,
quality=50,
differentiable=True).cuda()
# self.gaussian = Gaussian().cuda()
# self.dropout = Dropout(self.config,keep_ratio_range=(0.5,0.75)).cuda()
# self.resize = Resize().cuda()
# self.crop_layer = Crop((0.2, 0.5), (0.2, 0.5)).cuda()
self.noise_layers.append(self.jpeg_layer_80)
self.noise_layers.append(self.jpeg_layer_90)
self.noise_layers.append(self.jpeg_layer_70)
self.noise_layers.append(self.jpeg_layer_60)
self.noise_layers.append(self.jpeg_layer_50)
