I am currently doing this. I really don’t know much about pytorch so I would like as much help as possible. The error is in the end. I am also using a Raspberry Pi.
These are the codes:
main.py:
import torch.utils.data
import torchvision.datasets as datasets
from StyleCNN import *
from utils import *
# CUDA Configurations
dtype = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
# Content and style
style = image_loader("styles/starry_night.jpg").type(dtype)
#content = image_loader("contents/result.jpg").type(dtype)
#pastiche = image_loader("contents/result.jpg").type(dtype)
#pastiche.data = torch.randn(pastiche.data.size()).type(dtype)
num_epochs = 3
N = 4
def main():
style_cnn = StyleCNN(style)
# Contents
coco = datasets.ImageFolder(root='data', transform=loader)
content_loader = torch.utils.data.DataLoader(coco, batch_size=N, shuffle=True)
for epoch in range(num_epochs):
for i, content_batch in enumerate(content_loader):
iteration = epoch * i + i
content_loss, style_loss, pastiches = style_cnn.train(content_batch)
if i % 10 == 0:
print("Iteration: %d" % (iteration))
print("Content loss: %f" % (content_loss.data[0]))
print("Style loss: %f" % (style_loss.data[0]))
if i % 500 == 0:
path = "outputs/%d_" % (iteration)
paths = [path + str(n) + ".png" for n in range(N)]
save_images(pastiches, paths)
path = "outputs/content_%d_" % (iteration)
paths = [path + str(n) + ".png" for n in range(N)]
save_images(content_batch, paths)
style_cnn.save()
main()
StyleCNN.py:
import torch.optim as optim
import torchvision.models as models
from GramMatrix import *
class StyleCNN(object):
def __init__(self, style):
super(StyleCNN, self).__init__()
self.style = style
self.content_layers = ['conv_4']
self.style_layers = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
self.content_weight = 1
self.style_weight = 1000
self.transform_network = nn.Sequential(nn.ReflectionPad2d(40),
nn.Conv2d(3, 32, 9, stride=1, padding=4),
nn.Conv2d(32, 64, 3, stride=2, padding=1),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.Conv2d(128, 128, 3, stride=1, padding=0),
nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
nn.Conv2d(32, 3, 9, stride=1, padding=4),
)
self.gram = GramMatrix()
self.loss = nn.MSELoss()
self.optimizer = optim.Adam(self.transform_network.parameters(), lr=1e-3)
self.use_cuda = torch.cuda.is_available()
if self.use_cuda:
self.loss_network.cuda()
self.gram.cuda()
def train(self, content):
self.optimizer.zero_grad()
content = content
style = self.style
pastiche = self.transform_network.forward(content)
content_loss = 0
style_loss = 0
i = 1
not_inplace = lambda layer: nn.ReLU(inplace=False) if isinstance(layer, nn.ReLU) else layer
for layer in list(self.loss_network.features):
layer = not_inplace(layer)
if self.use_cuda:
layer.cuda()
pastiche, content, style = layer.forward(pastiche), layer.forward(content), layer.forward(style)
if isinstance(layer, nn.Conv2d):
name = "conv_" + str(i)
if name in self.content_layers:
content_loss += self.loss(pastiche * self.content_weight, content.detach() * self.content_weight)
if name in self.style_layers:
pastiche_g, style_g = self.gram.forward(pastiche), self.gram.forward(style)
style_loss += self.loss(pastiche_g * self.style_weight, style_g.detach() * self.style_weight)
if isinstance(layer, nn.ReLU):
i += 1
total_loss = content_loss + style_loss
total_loss.backward()
self.optimizer.step()
return self.pastiche
GramMatrix.py:
import torch
import torch.nn as nn
class GramMatrix(nn.Module):
def forward(self, input):
a, b, c, d = input.size()
features = input.view(a * b, c * d)
G = torch.mm(features, features.t())
return G.div(a * b * c * d)
This is the error:
Traceback (most recent call last):
File "main.py", line 47, in <module>
main()
File "main.py", line 30, in main
content_loss, style_loss, pastiches = style_cnn.train(content_batch)
File "/home/pi/Desktop/python_torch/StyleCNN.py", line 50, in train
pastiche = self.transform_network.forward(content)
File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/container.py", line 97, in forward
input = module(input)
File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/module.py", line 491, in __call__
result = self.forward(*input, **kwargs)
File "/usr/local/lib/python2.7/dist-packages/torch/nn/modules/padding.py", line 172, in forward
return F.pad(input, self.padding, 'reflect')
File "/usr/local/lib/python2.7/dist-packages/torch/nn/functional.py", line 2711, in pad
assert len(pad) // 2 <= input.dim(), 'Padding length too large'
AttributeError: 'list' object has no attribute 'dim'