Neural style transfer stack expects each tensor to be equal size

extremety1989: <p>Need help, it is simple neural style transfer script, it works by training but after saving and loading model state, i got this error</p>
<p>RuntimeError: stack expects each tensor to be equal size, but got [1, 64, 256, 256] at entry 0 and [1, 128, 128, 128] at entry 1<br>
at save_image(new_img, ‘output.png’)<br>
i do not understand what is wrong with it, i transformed the image size …</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.utils import save_image
device = torch.device("cuda" if torch.cuda.is_available else "cpu")

image_size  = 256

class VGG(nn.Module):

  def __init__(self):
    super(VGG,self).__init__()
    self.chosen_features = ["0","5","10","19","28"]
    self.model = models.vgg19(pretrained=True).features[:29]

  def forward(self,x):
    features = []
    for layer_num,layer  in enumerate(self.model):
      x = layer(x)
      if str(layer_num) in self.chosen_features:
        features.append(x)
    return features

def load_image(name):
  image = Image.open(name)
  image = loader(image).unsqueeze(0).to(device)
  return image


loader = transforms.Compose([
    transforms.Resize((image_size,image_size)),
    transforms.ToTensor(),
])
model = VGG().to(device).eval()

original_img = load_image("and.png")
style_img = load_image("style.jpg")
generated = original_img.clone().requires_grad_(True)
learning_rate = 0.001
alpha = 1
beta = 0.01
optimizer =  optim.Adam([generated],lr=learning_rate)

for step in range(200):
  optimizer.zero_grad()
  generated_features = model(generated)
  original_img_features = model(original_img)
  style_features = model(style_img)
  style_loss = original_loss = 0
  for gen_feature,orig_feature,style_feature in zip(
      generated_features,original_img_features,style_features
  ):
      batch_size,channel,height,width = gen_feature.shape
      original_loss += torch.mean((gen_feature - orig_feature)**2)
      G = gen_feature.view(channel,height * width) @ gen_feature.view(channel,height * width).T
      A = style_feature.view(channel,height * width) @ style_feature.view(channel,height * width).T
      style_loss += torch.mean((G - A)**2)
  total_loss = alpha * original_loss + beta * style_loss
  total_loss.backward()
  optimizer.step()
  if step % 10 == 0:
    print(total_loss)
    save_image(generated,"gen.png")#works
  torch.save({'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict()},
            "saved_model_for_styletransfer.pth")

  
with torch.no_grad(): 
  checkpoint = torch.load("saved_model_for_styletransfer.pth")
  model = VGG().to(device).eval()
  model.load_state_dict(checkpoint['model_state_dict'])
  new_img = model(load_image("and.png"))
  print(new_img)
  save_image(new_img, 'output.png')#fails
</code></pre>

ptrblck: <p>Your model returns a list of intermediate features in different shapes, so storing this list as a single image won’t work.<br>
You can check the shape of each tensor by iterating the output:</p>
<pre><code class="lang-python">for i in new_img:
    print(i.shape)
</code></pre>

extremety1989: <p>thanks for reply, well i got this<br>
torch.Size([1, 64, 256, 256])<br>
torch.Size([1, 128, 128, 128])<br>
torch.Size([1, 256, 64, 64])<br>
torch.Size([1, 512, 32, 32])<br>
torch.Size([1, 512, 16, 16])</p>
<p>i tried to save as eval and cpu:</p>
<pre><code class="lang-auto"># model.eval().cpu()          
# torch.save({"optimizer":copy.deepcopy(optimizer.state_dict()), 
#             "model":copy.deepcopy(model.state_dict())
#             },"saved_model_for_styletransfer_cpu.pth")

with torch.no_grad(): 
  checkpoint = torch.load("saved_model_for_styletransfer_cpu.pth")
  for k in list(checkpoint.keys()):
    if re.search(r'in\d+\.running_(mean|var)$', k):
        del checkpoint[k]
  model = VGG().cpu()
  model.load_state_dict(checkpoint['model'])
  ad = load_image("and.png").cpu()
  new_img = model(ad)
  for i in new_img:
    print(i.shape)
  save_image(new_img, 'output.png')
</code></pre>
<p>well training works, error is getting lower when i load model, but if i want only to do forward pass  i get that error… i do not understand why …</p>

ptrblck: <aside class="quote no-group" data-username="extremety1989" data-post="3" data-topic="134933">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/extremety1989/48/37103_2.png" class="avatar"> extremety1989:</div>
<blockquote>
<p>but if i want only to do forward pass i get that error… i do not understand why …</p>
</blockquote>
</aside>
<p>The forward pass is not raising the error, but the <code>save_image</code> applied on a list of tensors with different shapes.<br>
Remove</p>
<pre data-code-wrap="python"><code class="lang-python">save_image(new_img, 'output.png')
</code></pre>
<p>and it should work or alternatively store each slice of each intermediate activation separately.</p>

extremety1989: <p>it is because im using optimizer =  optim.Adam([generated],lr=learning_rate)<br>
optimizing the image and not model parameters , but i do not know the way how to get work without training, just to style image</p>

ptrblck: <p>I’m not sure I understand the explanation. The code fails in saving the <code>list</code> of tensors so how does it relate to the model training and optimizer?</p>