demo.py

import torchvision.models as models
from PIL import Image
from model import Generator
import torchvision.transforms as transforms
import torch.nn.functional as F
import torch.optim as optim
import torch.nn as nn
import torch
import os
os.chdir('/data1/ycjh/reb/pixel2style2pixel')
print('Imported')

ver = ''
os.makedirs('test'+ver, exist_ok=True)
os.makedirs('mine'+ver, exist_ok=True)

# 参数声明
batch_size = 32
epochs = 10000
log_dir = 'mine'+ver+'/mine.pth'  # 模型保存路径
g_ema = Generator(
    1024, 512, 8, channel_multiplier=2
).to('cuda')
checkpoint = torch.load('stylegan2-ffhq-config-f.pt')

g_ema.load_state_dict(checkpoint["g_ema"])
mean_latent = None

g_ema.eval()

# ln = [64, 64, 64, 64, 32, 32, 32, 16, 16, 16, 16, 8, 8, 8, 8, 4, 4, 2, 2]
# ch = [192//ln[i] for i in range(19)]


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.mp = nn.ModuleList()
        self.li = nn.ModuleList()
        for _ in range(1, 19):
            resnet18 = models.resnet18(pretrained=True)
            res = nn.Sequential(*(list(resnet18.children())[0:8]))
            self.li.append(nn.Linear(2048, 512))
            self.mp.append(res)

    def forward(self, x, N=-1, lat=None):
        N = max(N, -1)
        vec = torch.zeros(x.shape[0], 18, 512).cuda()
        if N != -1:
            vec = lat.clone()
            y = self.mp[N](x.clone())
            y = y.view(y.size(0), -1)
            y = self.li[N](y)
            y = F.leaky_relu(y)
            y = y.view(y.shape[0], 512)
            vec[:, 17 - N, :] = y
        else:
            for i in range(18):
                y = self.mp[i](x.clone())
                y = y.view(y.size(0), -1)
                y = self.li[i](y)
                y = F.leaky_relu(y)
                y = y.view(y.shape[0], 512)
                vec[:, 17 - i, :] = y
        x, _ = g_ema(
            [vec], input_is_latent=True, randomize_noise=False
        )
        x = x.mul(255).add(0.5).clamp(0, 255).floor()/255
        return x, vec


model = Net().cuda()
criterion = nn.L1Loss()
optimizer = optim.SGD(model.parameters(), lr=0.1)


logs = open('log'+ver+'.txt', 'a')
# 模型训练


def train(model, batch_s, ep):
    global lat
    model.train()
    train_loss = 0
    lo = [0]*18
    for i in range(batch_s):
        # N=max(0,17-ep//500)
        for j in range(18):
            N = j % 18
            vec = torch.randn(6, 512, device='cuda').cuda()
            with torch.no_grad():
                y, lat = g_ema(
                    [vec], return_latents=True, randomize_noise=False
                )
                y = y.mul(255).add(0.5).clamp(0, 255).floor()/255
                sv = F.interpolate(y, (64, 64), mode='bilinear',
                                   align_corners=False)
                sv = sv.mul(255).add(0.5).clamp(0, 255).floor()/255
            optimizer.zero_grad()
            y_hat, y_lat = model(sv, N, lat)
            loss = criterion(y_hat, y)
            # loss = criterion(y_hat, y)+criterion(lat, y_lat)
            # loss = criterion(lat, y_lat)
            loss.backward()
            optimizer.step()
            train_loss += loss.item()
            lo[N] += loss.item()
            print('\r%d%% %d %.4f %.4f' %
                  (int((i+(j+1)/18)/batch_s*100), N, loss.item(), train_loss/(i+(j+1)/18)), end='')
        if i+1 % 16 == 0:
            print()
    loss_mean = train_loss / (i+1)
    S = 'Train Epoch: {:0>5d}\t Loss: {:.6f}\t'.format(ep, loss_mean)
    print('\n'+S)
    logs.write(S+'\t')
    for i in lo:
        logs.write('\t%.3f' % i)
    logs.write('\n')
    logs.flush()

# 模型测试


def test(model, fr=-20, to=18):
    global lat, logs, N
    model.eval()
    test_loss = 0
    with torch.no_grad():
        for i in range(fr, to):
            N = i
            x = torch.randn(1, 512, device='cuda')
            y, lat = g_ema(
                [x], return_latents=True
            )
            y = y.mul(255).add(0.5).clamp(0, 255).floor()/255
            x = x.cuda()
            y = y.cuda()
            x = F.interpolate(y, (64, 64))
            x = x.mul(255).add(0.5).clamp(0, 255).floor()/255
            if N < 0:
                y_hat, _ = model(x, N)
            else:
                y_hat, _ = model(x, N, lat)
            test_loss += criterion(F.interpolate(y_hat,
                                   (64, 64), mode='bilinear', align_corners=False), x).item()
            y_ = Image.fromarray(y_hat[0].mul(255).add_(0.5).clamp_(0, 255).permute(
                1, 2, 0).to('cpu', torch.uint8).numpy())
            z_ = Image.fromarray(y[0].mul(255).add_(0.5).clamp_(0, 255).permute(
                1, 2, 0).to('cpu', torch.uint8).numpy())
            dic = 'test'+ver+'/'
            y_.save((dic+str(i).zfill(3)+'res.jpg').replace('-', '@'))
            y_.resize((64, 64)).save(
                (dic+str(i).zfill(3)+'smr.jpg').replace('-', '@'))
            z_.save((dic+str(i).zfill(3)+'ori.jpg').replace('-', '@'))
            z_.resize((64, 64)).save(
                (dic+str(i).zfill(3)+'smo.jpg').replace('-', '@'))


def main():

    # 如果有保存的模型，则加载模型，并在其基础上继续训练
    if os.path.exists(log_dir):
        checkpoint = torch.load(log_dir)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        start_epoch = checkpoint['epoch']
        print('加载 epoch {} 成功！'.format(start_epoch))
    else:
        start_epoch = 0
        print('无保存模型，将从头开始训练！')
    # 如果test_flag=True,则加载已保存的模型
    test_flag = True  # 测试标志，True时加载保存好的模型进行测试
    if False:
        test(model, -1, 0)
        exit(0)
    if test_flag:
        model.eval()
        # 加载保存的模型直接进行测试机验证，不进行此模块以后的步骤
        with torch.no_grad():
            while True:
                s = input('File Name:')
                y = Image.open(s)
                y = transforms.functional.to_tensor(y).cuda()
                y = y.unsqueeze(0)
                # print(y.shape, y.max(), y.min(), y.mean(), y.std(), y.dtype)
                x = F.interpolate(y, (64, 64))
                y_hat, _ = model(x, -1)
                y_ = Image.fromarray(y_hat[0].mul(255).add_(0.5).clamp_(0, 255).permute(
                    1, 2, 0).to('cpu', torch.uint8).numpy())
                y_.save('cov-'+s)
                # break
        return
    for epoch in range(start_epoch+1, epochs):
        train(model,  16, epoch)
        test(model)
        # 保存模型
        state = {'model': model.state_dict(
        ), 'optimizer': optimizer.state_dict(), 'epoch': epoch}
        torch.save(state, log_dir)


if __name__ == '__main__':
    main()