# This is the code from the p1ch2/3_cyclegan notebook
import torch
import torch.nn as nn
class ResNetBlock(nn.Module):
def __init__(self, dim):
super(ResNetBlock, self).__init__()
self.conv_block = self.build_conv_block(dim)
def build_conv_block(self, dim):
conv_block = []
conv_block += [nn.ReflectionPad2d(1)]
conv_block += [
nn.Conv2d(dim, dim, kernel_size=3, padding=0, bias=True),
nn.InstanceNorm2d(dim),
nn.ReLU(True),
]
conv_block += [nn.ReflectionPad2d(1)]
conv_block += [
nn.Conv2d(dim, dim, kernel_size=3, padding=0, bias=True),
nn.InstanceNorm2d(dim),
]
return nn.Sequential(*conv_block)
def forward(self, x):
out = x + self.conv_block(x)
return out
class ResNetGenerator(nn.Module):
def __init__(self, input_nc=3, output_nc=3, ngf=64, n_blocks=9):
assert n_blocks >= 0
super(ResNetGenerator, self).__init__()
self.input_nc = input_nc
self.output_nc = output_nc
self.ngf = ngf
model = [
nn.ReflectionPad2d(3),
nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=True),
nn.InstanceNorm2d(ngf),
nn.ReLU(True),
]
n_downsampling = 2
for i in range(n_downsampling):
mult = 2**i
model += [
nn.Conv2d(
ngf * mult,
ngf * mult * 2,
kernel_size=3,
stride=2,
padding=1,
bias=True,
),
nn.InstanceNorm2d(ngf * mult * 2),
nn.ReLU(True),
]
mult = 2**n_downsampling
for i in range(n_blocks):
model += [ResNetBlock(ngf * mult)]
for i in range(n_downsampling):
mult = 2 ** (n_downsampling - i)
model += [
nn.ConvTranspose2d(
ngf * mult,
int(ngf * mult / 2),
kernel_size=3,
stride=2,
padding=1,
output_padding=1,
bias=True,
),
nn.InstanceNorm2d(int(ngf * mult / 2)),
nn.ReLU(True),
]
model += [nn.ReflectionPad2d(3)]
model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]
model += [nn.Tanh()]
self.model = nn.Sequential(*model)
def forward(self, input):
# here we move to 0-1 input and 0-1 output
# usually one would think about writing this differently
# for efficiency (e.g. absorbing the 255 into the first conv
return self.model(input * 255) / 2 + 0.5
def get_pretrained_model(model_path, map_location=None):
netG = ResNetGenerator()
model_data = torch.load(model_path, map_location=map_location)
netG.load_state_dict(model_data)
netG.eval()
for p in netG.parameters():
netG.requires_grad_(False)
return netG
if __name__ == "__main__":
import sys
if len(sys.argv) < 3:
print("Call as {} zebra_weights.pt traced_zebra_model.pt".format(sys.argv[0]))
sys.exit(1)
model = get_pretrained_model(sys.argv[1], map_location="cpu")
traced_model = torch.jit.trace(model, torch.randn(1, 3, 227, 227))
traced_model.save(sys.argv[2])
# img = Image.open("../data/p1ch2/horse.jpg")
# out_img.save('../data/p1ch2/zebra.jpg')