PRSRGAN test

pr_vgg.py

import torch
import torch.nn as nn
import torchvision.models as models
from pprint import pprint

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        model = models.vgg19(pretrained=False)
        
        pprint(list(model.children()))
        
        model = list(model.children())[:-2]
        model = model[0]
        model = list(model.children())
        
        vgg_arr = []
        
        for xx in model:
            if 'Max' not in  xx.__class__.__name__:
                vgg_arr.append(xx)
                
        self.net = nn.Sequential(*vgg_arr)
        
    def forward(self, x):
        x = self.net(x)
        
        return x


device = torch.device("cuda:0")

model = Net().to(device)
model.eval()

t = torch.randn((1, 3, 224, 224)).to(device)

out = model(t)

print(out.shape)