# Style loss class GramMatrix(nn.Module): def forward(self, input): a, b, c, d = input.size() # a=batch size(=1) # b=number of feature maps # (c,d)=dimensions of a f. map (N=c*d) features = input.view(a * b, c * d) # resise F_XL into \hat F_XL G = torch.mm(features, features.t()) # compute the gram product # we 'normalize' the values of the gram matrix # by dividing by the number of element in each feature maps. return G.div(a * b * c * d) class StyleLoss(nn.Module): def __init__(self, target, weight): super(StyleLoss, self).__init__() self.target = target.detach() * weight self.weight = weight self.gram = GramMatrix() self.criterion = nn.MSELoss() def forward(self, input): self.output = input.clone() self.G = self.gram(input) self.G.mul_(self.weight) self.loss = self.criterion(self.G, self.target) return self.output def backward(self, retain_graph=True): self.loss.backward(retain_graph=retain_graph) return self.loss

这段代码是用于计算风格损失的。其中，GramMatrix类用于计算输入的Gram矩阵，即特征图的协方差矩阵，以表达输入的风格信息；StyleLoss类则用于计算输入与目标风格之间的均方误差，作为风格损失。

在forward方法中，输入被克隆为输出，并使用GramMatrix计算出输入的Gram矩阵，再乘以权重，最后计算均方误差得到风格损失。在backward方法中，反向传播损失，并返回损失值。