Изменения
→Пример кода на PyTorch
<math>L_{style}(S, G) = \sum\limits_{l=0}^L w_l * L_{GM}(S, G, l)</math>
== Пример кода на PyTorch Python ==
'''Content Loss'''
class ContentLoss(nn.Module):
'''Style Loss''' def gram_matrix(input): a, b, c, d = input.size() # a=batch size(=1) <font color="green"># b=number of feature maps</font> <font color="green"># (c,d)=dimensions of a f. map (N=c*d)</font> features = input.view(a * b, c * d) # resise F_XL into \hat F_XL G = torch.mm(features, features.t()) # compute the gram product <font color="green"># we 'normalize' the values of the gram matrix</font> <font color="green"># by dividing by the number of element in each feature maps.</font> return G.div(a * b * c * d)
class StyleLoss(nn.Module):
'''Importing the Model'''
cnn = models.vgg19(pretrained=True).features.to(device).eval()
'''Normalization'''
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)
<font color="green"># create a module to normalize input image so we can easily put it in a</font>
<font color="green"># nn.Sequential</font>
class Normalization(nn.Module):
def __init__(self, mean, std):
super(Normalization, self).__init__()
<font color="green"># .view the mean and std to make them [C x 1 x 1] so that they can</font>
<font color="green"># directly work with image Tensor of shape [B x C x H x W].</font>
<font color="green"># B is batch size. C is number of channels. H is height and W is width.</font>
self.mean = torch.tensor(mean).view(-1, 1, 1)
self.std = torch.tensor(std).view(-1, 1, 1)
def forward(self, img):
<font color="green"># normalize img</font>
return (img - self.mean) / self.std
'''Adding our content loss and style loss layers'''
<font color="green"># desired depth layers to compute style/content losses :</font>
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, normalization_mean, normalization_std,
style_img, content_img,
content_layers=content_layers_default,
style_layers=style_layers_default):
cnn = copy.deepcopy(cnn)
<font color="green"># normalization module</font>
normalization = Normalization(normalization_mean, normalization_std).to(device)
<font color="green"># just in order to have an iterable access to or list of content/style losses</font>
content_losses = []
style_losses = []
<font color="green"># assuming that cnn is a nn.Sequential, so we make a new nn.Sequential</font>
<font color="green"># to put in modules that are supposed to be activated sequentially</font>
model = nn.Sequential(normalization)
i = 0 <font color="green"># increment every time we see a conv</font>
for layer in cnn.children():
if isinstance(layer, nn.Conv2d):
i += 1
name = 'conv_{}'.format(i)
elif isinstance(layer, nn.ReLU):
name = 'relu_{}'.format(i)
<font color="green"># The in-place version doesn't play very nicely with the ContentLoss</font>
<font color="green"># and StyleLoss we insert below. So we replace with out-of-place</font>
<font color="green"># ones here.</font>
layer = nn.ReLU(inplace=False)
elif isinstance(layer, nn.MaxPool2d):
name = 'pool_{}'.format(i)
elif isinstance(layer, nn.BatchNorm2d):
name = 'bn_{}'.format(i)
else:
raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))
model.add_module(name, layer)
if name in content_layers:
<font color="green"># add content loss:</font>
target = model(content_img).detach()
content_loss = ContentLoss(target)
model.add_module("content_loss_{}".format(i), content_loss)
content_losses.append(content_loss)
if name in style_layers:
<font color="green"># add style loss:</font>
target_feature = model(style_img).detach()
style_loss = StyleLoss(target_feature)
model.add_module("style_loss_{}".format(i), style_loss)
style_losses.append(style_loss)
==См. также==