Kaggle competition: Generative Dog Images
from typing import List, Tuple
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
import torch.utils.data
import torchvision
from torchvision import transforms, datasets
from torchvision.utils import save_image
from tqdm import tqdm_notebook as tqdm<init_hyperparams>
writer = tbx.SummaryWriter("." )
EPOCHS = 270
BATCH_SIZE = 64
NORMALIZE_Z = True
Z_SIZE = 128
N_FEATURES = 32
CONDITIONAL = True
SELF_ATTENTION = True
SPECTRAL_IN_G = True
SPECTRAL_IN_D = True
BN_IN_D = False
RESIDUAL_IN_D = True
GENERATOR_ACTIVATION = F.relu
DISCRIMINATOR_ACTIVATION = F.relu
class LossType(enum.Enum):
LS = 1
RaLS = 2
Hinge = 3
RaHinge = 4
def requires_real_for_generator(self)-> bool:
return self == LossType.RaLS
LOSS = LossType.Hinge
BETA1 = 0.0
BETA2 = 0.999
LR_GENERATOR = 0.0001
LR_DISCRIMINATOR = 0.0004
CROP_PROB = 1.0
SHOW_FREQ = None
SHOW_FREQ_EPOCH = EPOCHS // 10<data_type_conversions>
start_time = time.time()
random.seed(0)
 np.random.seed(0)
 torch.manual_seed(0 )
def sample(size):
z = np.random.normal(loc=0, scale=1, size=(size, Z_SIZE)).astype(np.float32)
 if NORMALIZE_Z:
z /= np.sqrt(np.sum(z ** 2, axis=1, keepdims=True))
return torch.as_tensor(z)
 torch.manual_seed(42)
 FIXED_NOISES = sample(16 )<define_variables>
def doggo_loader(path):
img = torchvision.datasets.folder.default_loader(path)
 annotation_basename = os.path.splitext(os.path.basename(path)) [0]
annotation_dirname = next(dirname for dirname in os.listdir('.. /input/annotation/Annotation/')if dirname.startswith(annotation_basename.split('_')[0]))
annotation_filename = os.path.join('.. /input/annotation/Annotation', annotation_dirname, annotation_basename)
 tree = ET.parse(annotation_filename)
 root = tree.getroot()
objects = root.findall('object')
 for idx, o in enumerate(objects):
bndbox = o.find('bndbox')
 xmin = int(bndbox.find('xmin' ).text)
 ymin = int(bndbox.find('ymin' ).text)
 xmax = int(bndbox.find('xmax' ).text)
 ymax = int(bndbox.find('ymax' ).text)
 w = xmax - xmin
h = ymax - ymin
l = max([w,h])
 if w > h:
ymin = ymin - int(( l-h)/2)
 else:
xmin = xmin - int(( l-w)/2)
 bbox =(xmin, ymin , xmin+l, ymin+l)
 c_img = img.crop(bbox ).copy()
c_img.thumbnail(( 100, 100), Image.ANTIALIAS)
 c_img.save('.. /class_images/train/image_{}_{}.png'.format(annotation_basename, idx))
return img
dataset = torchvision.datasets.ImageFolder(
 '.. /input/all-dogs/',
loader=doggo_loader,
transform=torchvision.transforms.Compose([
transforms.RandomResizedCrop(64, scale=(1.0, 1.0), ratio=(1., 1.)),
torchvision.transforms.ToTensor() ,
])
 )
def extract_class(path: str)-> str:
name = Path(path ).name
idx = name.index("_")
 return name[:idx]
def gather_image_and_annots(annot_path: str = ".. /input/annotation/Annotation/", image_path: str = ".. /input/all-dogs/all-dogs/")-> List[Tuple[str, str]]:
results: List[Tuple[str, str]] = []
breeds = os.listdir(annot_path)
 for breed in breeds:
for dog in os.listdir(os.path.join(annot_path, breed)) :
annot = os.path.join(annot_path, breed, dog)
 image = os.path.join(image_path, dog + ".jpg")
 if os.path.exists(image):
results.append(( image, annot))
return results
def get_classes(path: str):
classes = set()
for filepath in Path(path ).glob("*.jpg"):
classes.add(extract_class(filepath))
return sorted(classes)
 def filter_small_boxes(boxes: List[List[int]])-> List[List[int]]:
return [x for x in boxes if min(x[2] - x[0], x[3] - x[1])>= 64]
def get_bounding_boxes(path: str, filter_small: bool = True)-> List[List[int]]:
tree = ET.parse(path)
 root = tree.getroot()
objects = root.findall('object')
 boxes: List[List[int]] = []
for o in objects:
bndbox = o.find('bndbox')
 xmin = int(bndbox.find('xmin' ).text)
 ymin = int(bndbox.find('ymin' ).text)
 xmax = int(bndbox.find('xmax' ).text)
 ymax = int(bndbox.find('ymax' ).text)
 boxes.append([xmin, ymin, xmax, ymax])
 if filter_small:
boxes = filter_small_boxes(boxes)
 return boxes
Boxes = List[List[int]]
def append_bounding_boxes(pairs: List[Tuple[str, str]], q: float = 0.99)-> List[Tuple[str, Boxes]]:
n = len(pairs)
 sizes = []
for img, _ in tqdm(pairs):
sizes.append(os.stat(img ).st_size)
 pairs = np.array(pairs)
 sizes = np.array(sizes)
 indices = sizes < np.quantile(sizes, q)
 pairs = pairs[indices]
results = []
for i,(img_path, annot)in enumerate(tqdm(pairs)) :
img = Image.open(img_path)
 if min(*img.size)< 64:
continue
boxes = get_bounding_boxes(annot, filter_small=True)
 if not boxes:
continue
results.append(( img_path, boxes))
print("Valid sized images: {} / {}".format(len(results), n))
return results
CLASSES = get_classes(".. /input/all-dogs/all-dogs/")
 PAIRS = append_bounding_boxes(gather_image_and_annots())
 IMAGES, BOXES = [p[0] for p in PAIRS], [p[1] for p in PAIRS]
len(CLASSES), len(PAIRS )<categorify>
if not os.path.exists('.. /class_images'):
os.mkdir('.. /class_images')
 if not os.path.exists('.. /class_images/train'):
os.mkdir('.. /class_images/train')
 dataloader = torch.utils.data.DataLoader(dataset, batch_size=50,
shuffle=False, num_workers=8)
 for i, data in enumerate(tqdm(dataloader)) :
pass
class Dataset(torch.utils.data.Dataset):
def __init__(self, image_files, boxes, transform, crop_prob: float = 0.5, jitter_ratio: float = 0.05):
super().__init__()
self.transform = transform
self.image_files = image_files
self.boxes = boxes
self.crop_prob = crop_prob
self.jitter_ratio = jitter_ratio
def __len__(self)-> int:
return len(self.image_files)
 def __getitem__(self, idx):
filepath = self.image_files[idx]
image = Image.open(str(filepath)).convert("RGB")
 if self.crop_prob >= 1.0 or np.random.random() < self.crop_prob:
boxes = self.boxes[idx]
image = self._crop_bounding_boxes(image, boxes)
 if CONDITIONAL:
class_id = CLASSES.index(extract_class(filepath))
else:
class_id = 0
image_tensor = self.transform(image)
 return image_tensor, torch.as_tensor(class_id)
 def _crop_bounding_boxes(self, image: Image.Image, boxes: List[List[int]])-> Image.Image:
i = np.random.randint(0, len(boxes))
xmin, ymin, xmax, ymax = boxes[i]
if min(xmax - xmin, ymax - ymin)< 64:
return image
x_center =(xmin + xmax)// 2
y_center =(ymin + ymax)// 2
size = max(xmax - xmin, ymax - ymin)
 if int(size * self.jitter_ratio):
x_jitter = np.random.randint(int(-size * self.jitter_ratio), int(size * self.jitter_ratio))
x_center += x_jitter
y_jitter = np.random.randint(int(-size * self.jitter_ratio), int(size * self.jitter_ratio))
y_center += y_jitter
size_jitter = np.random.randint(0, int(size * self.jitter_ratio))
xmin = max(0, x_center -(size + size_jitter)// 2)
 ymin = max(0, y_center -(size + size_jitter)// 2)
 xmax = min(image.size[0], x_center +(size + size_jitter)// 2)
 ymax = min(image.size[1], y_center +(size + size_jitter)// 2)
 return image.crop(( xmin, ymin, xmax, ymax))<create_dataframe>
workers = 8
batch_size=32
nz = 128
nch_g = 64
nch_d = 64
n_epoch = 20000000
lr = 0.0002
beta1 = 0.5
outf = './result_lsgan'
display_interval = 100
save_fake_image_interval = 1500
plt.rcParams['figure.figsize'] = 10, 6
try:
os.makedirs(outf, exist_ok=True)
 except OSError as error:
print(error)
 pass
random.seed(0)
 np.random.seed(0)
 torch.manual_seed(0)
 def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv')!= -1:
m.weight.data.normal_(0.0, 0.02)
 m.bias.data.fill_(0)
 elif classname.find('Linear')!= -1:
m.weight.data.normal_(0.0, 0.02)
 m.bias.data.fill_(0)
 elif classname.find('BatchNorm')!= -1:
m.weight.data.normal_(1.0, 0.02)
 m.bias.data.fill_(0 )
transform = transforms.Compose([
transforms.Resize(64),
transforms.RandomCrop(64),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor() ,
transforms.Normalize(( 0.5, 0.5, 0.5),(0.5, 0.5, 0.5)) ,
])
 dataset = Dataset(IMAGES, BOXES, transform, crop_prob=CROP_PROB)
 dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, batch_size=BATCH_SIZE, num_workers=4)
 print(len(dataset), len(dataloader))
_ = next(iter(dataloader))<init_hyperparams>
def init_linear(linear):
init.xavier_uniform_(linear.weight)
 linear.bias.data.zero_()
def init_conv(conv, glu=True):
init.xavier_uniform_(conv.weight)
 if conv.bias is not None:
conv.bias.data.zero_()
def spectral_init(module, use_norm: bool, gain=1):
init.kaiming_uniform_(module.weight, gain)
 if module.bias is not None:
module.bias.data.zero_()
if use_norm:
return spectral_norm(module)
 else:
return module
def leaky_relu(input):
return F.leaky_relu(input, negative_slope=0.2)
 class SelfAttention(nn.Module):
def __init__(self, in_channel: int, use_norm: bool, gain=1):
super().__init__()
self.query = spectral_init(nn.Conv1d(in_channel, in_channel // 8, 1),
use_norm=use_norm,
gain=gain)
 self.key = spectral_init(nn.Conv1d(in_channel, in_channel // 8, 1),
use_norm=use_norm,
gain=gain)
 self.value = spectral_init(nn.Conv1d(in_channel, in_channel, 1),
use_norm=use_norm,
gain=gain)
 self.gamma = nn.Parameter(torch.tensor(0.0))
def forward(self, input):
shape = input.shape
flatten = input.view(shape[0], shape[1], -1)
 query = self.query(flatten ).permute(0, 2, 1)
 key = self.key(flatten)
 value = self.value(flatten)
 query_key = torch.bmm(query, key)
 attn = F.softmax(query_key, 1)
 attn = torch.bmm(value, attn)
 attn = attn.view(*shape)
 out = self.gamma * attn + input
return out
class ConditionalNorm(nn.Module):
def __init__(self, in_channel, n_class):
super().__init__()
self.bn = nn.BatchNorm2d(in_channel, affine=False)
 self.embed = nn.Embedding(n_class, in_channel * 2)
 self.embed.weight.data[:, :in_channel] = 1
self.embed.weight.data[:, in_channel:] = 0
def forward(self, input, class_id):
out = self.bn(input)
 embed = self.embed(class_id)
 gamma, beta = embed.chunk(2, 1)
 gamma = gamma.unsqueeze(2 ).unsqueeze(3)
 beta = beta.unsqueeze(2 ).unsqueeze(3)
 out = gamma * out + beta
return out
def snconv2d(
 in_features: int,
out_features: int,
kernel_size: int,
stride: int,
padding: int,
use_norm: bool, 
)-> nn.Module:
conv = nn.Conv2d(in_features, out_features, kernel_size, stride=stride, padding=padding, bias=True)
 init.xavier_uniform_(conv.weight)
 conv.bias.data.zero_()
if use_norm:
return spectral_norm(conv)
 else:
return conv
def snlinear(in_features: int, out_features: int, use_norm: bool)-> nn.Module:
linear = nn.Linear(in_features, out_features, bias=True)
 init.xavier_uniform_(linear.weight)
 linear.bias.data.zero_()
if use_norm:
return spectral_norm(linear)
 else:
return linear
def sn_embedding(num_classes: int, embedding_dim: int, use_norm: bool)-> nn.Module:
emb = nn.Embedding(num_classes, embedding_dim)
 init.xavier_uniform_(emb.weight)
 if use_norm:
return spectral_norm(emb)
 else:
return emb<categorify>
class PixelwiseNorm(nn.Module):
def __init__(self):
super(PixelwiseNorm, self ).__init__()
def forward(self, x, alpha=1e-8):

 y = x.pow(2.).mean(dim=1, keepdim=True ).add(alpha ).sqrt()
y = x / y
return y
class MinibatchStdDev(torch.nn.Module):

 def __init__(self):

 super(MinibatchStdDev, self ).__init__()
def forward(self, x, alpha=1e-8):

 batch_size, _, height, width = x.shape
y = x - x.mean(dim=0, keepdim=True)
 y = torch.sqrt(y.pow(2.).mean(dim=0, keepdim=False)+ alpha)
 y = y.mean().view(1, 1, 1, 1)
 y = y.repeat(batch_size,1, height, width)
 y = torch.cat([x, y], 1)
 return y
class Generator(nn.Module):
def __init__(self, nz, nfeats, nchannels):
super(Generator, self ).__init__()
self.conv1 = spectral_norm(nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False))
self.conv2 = spectral_norm(nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False))
self.conv3 = spectral_norm(nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False))
self.conv4 = spectral_norm(nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False))
self.conv5 = spectral_norm(nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False))
self.conv6 = spectral_norm(nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))
self.pixnorm = PixelwiseNorm()
def forward(self, x):
x = F.leaky_relu(self.conv1(x))
x = F.leaky_relu(self.conv2(x))
x = self.pixnorm(x)
 x = F.leaky_relu(self.conv3(x))
x = self.pixnorm(x)
 x = F.leaky_relu(self.conv4(x))
x = self.pixnorm(x)
 x = F.leaky_relu(self.conv5(x))
x = self.pixnorm(x)
 x = torch.tanh(self.conv6(x))
return x
def denorm(self, z):
z = self.forward(z)
 z = z * 0.5 + 0.5
return z
class Discriminator(nn.Module):
def __init__(self, nchannels, nfeats):
super(Discriminator, self ).__init__()
self.conv1 = nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)
 self.conv2 = spectral_norm(nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))
self.bn2 = nn.BatchNorm2d(nfeats * 2)
 self.conv3 = spectral_norm(nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False))
self.bn3 = nn.BatchNorm2d(nfeats * 4)
 self.conv4 = spectral_norm(nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False))
self.bn4 = nn.MaxPool2d(2)
 self.batch_discriminator = MinibatchStdDev()
self.pixnorm = PixelwiseNorm()
self.conv5 = spectral_norm(nn.Conv2d(nfeats * 8 +1, 1, 2, 1, 0, bias=False))
def forward(self, x):
x = F.leaky_relu(self.conv1(x), 0.2)
 x = F.leaky_relu(self.bn2(self.conv2(x)) , 0.2)
 x = F.leaky_relu(self.bn3(self.conv3(x)) , 0.2)
 x = F.leaky_relu(self.bn4(self.conv4(x)) , 0.2)
 x = self.batch_discriminator(x)
 x = torch.sigmoid(self.conv5(x))
return x.view(-1, 1)
def upsample_conv(in_features: int, out_features: int, kernel_size: int, stride: int, padding: int)-> nn.Module:
return nn.Sequential(
 nn.Upsample(scale_factor=2),
snconv2d(in_features, out_features, kernel_size=kernel_size, stride=stride, padding=padding, use_norm=SPECTRAL_IN_G), 
)
 class GeneratorBlock(nn.Module):
def __init__(self, in_features: int, out_features: int, num_classes: int)-> None:
super().__init__()
self.in_features = in_features
if CONDITIONAL:
self.bn0 = ConditionalNorm(in_features, num_classes)
 self.bn1 = ConditionalNorm(out_features, num_classes)
 else:
self.bn0 = nn.BatchNorm2d(in_features)
 self.bn1 = nn.BatchNorm2d(out_features)
 self.upconv = upsample_conv(in_features, out_features, kernel_size=3, stride=1, padding=1)
 self.conv1 = snconv2d(out_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_G)
 self.skip_upconv = upsample_conv(in_features, out_features, kernel_size=1, stride=1, padding=0)
 self.activation = GENERATOR_ACTIVATION
def forward(self, x, class_id):
x0 = x
if CONDITIONAL:
x = self.activation(self.bn0(x, class_id))
else:
x = self.activation(self.bn0(x))
x = self.upconv(x)
 if CONDITIONAL:
x = self.activation(self.bn1(x, class_id))
else:
x = self.activation(self.bn1(x))
x = self.conv1(x)
 x0 = self.skip_upconv(x0)
 return x + x0
class Generator(nn.Module):
def __init__(self, z_size: int = Z_SIZE, num_classes: int = len(CLASSES), features: int = N_FEATURES)-> None:
super().__init__()
self.linear0 = snlinear(Z_SIZE, features * 4 * 4 * 8, use_norm=SPECTRAL_IN_G)
 self.block0 = GeneratorBlock(features * 8, features * 8, num_classes=num_classes)
 self.block1 = GeneratorBlock(features * 8, features * 4, num_classes=num_classes)
 self.block2 = GeneratorBlock(features * 4, features * 2, num_classes=num_classes)
 if SELF_ATTENTION:
self.attn = SelfAttention(features * 2, use_norm=SPECTRAL_IN_G)
 self.block3 = GeneratorBlock(features * 2, features * 1, num_classes=num_classes)
 self.bn = nn.BatchNorm2d(features * 1)
 self.last_conv = snconv2d(features, 3, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_G)
 self.activation = GENERATOR_ACTIVATION
def forward(self, z, class_id):
x = self.linear0(z)
 x = x.view(x.size(0), -1, 4, 4)
 x = self.block0(x, class_id)
 x = self.block1(x, class_id)
 x = self.block2(x, class_id)
 if SELF_ATTENTION:
x = self.attn(x)
 x = self.block3(x, class_id)
 x = self.activation(self.bn(x))
x = self.last_conv(x)
 x = torch.tanh(x)
 return x
generator = Generator()
print(generator(torch.randn(1, Z_SIZE), torch.zeros(1, dtype=torch.long)).shape )<normalization>
class DogDataset(Dataset):

 def __init__(self, data_path, transform=None):

 self.data_list = sorted(glob.glob(data_path, recursive=True))
self.transform = transform
def __len__(self):
return len(self.data_list)
 def __getitem__(self, idx):
image = Image.open(self.data_list[idx])
 if self.transform:
image = self.transform(image)
 return image
def downsample(x):
return F.avg_pool2d(x, kernel_size=2, stride=2)
 class DiscriminatorBlock(nn.Module):
def __init__(self, in_features: int, out_features: int, down=True, activation=DISCRIMINATOR_ACTIVATION):
super().__init__()
self.down = down
self.activation = activation
if BN_IN_D:
self.bn0 = nn.BatchNorm2d(in_features)
 self.bn1 = nn.BatchNorm2d(out_features)
 self.conv0 = snconv2d(in_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)
 self.conv1 = snconv2d(out_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)
 if RESIDUAL_IN_D and(down or in_features != out_features):
self.skip_conv = snconv2d(in_features, out_features, kernel_size=1, stride=1, padding=0, use_norm=SPECTRAL_IN_D)
 else:
self.skip_conv = None
def forward(self, x):
x0 = x
if BN_IN_D:
x = self.bn0(x)
 x = self.activation(x)
 x = self.conv0(x)
 if BN_IN_D:
x = self.bn1(x)
 x = self.activation(x)
 x = self.conv1(x)
 if self.down:
x = downsample(x)
 if self.skip_conv is not None:
x0 = self.skip_conv(x0)
 if self.down:
x0 = downsample(x0)
 if RESIDUAL_IN_D:
x = x0 + x
return x
class DiscriminatorDownOptimizedBlock(nn.Module):
def __init__(self, in_features: int, out_features: int, activation=DISCRIMINATOR_ACTIVATION)-> None:
super().__init__()
self.conv1 = snconv2d(in_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)
 self.activation = activation
self.conv2 = snconv2d(out_features, out_features, kernel_size=3, stride=1, padding=1, use_norm=SPECTRAL_IN_D)
 if RESIDUAL_IN_D:
self.skip_conv = snconv2d(in_features, out_features, kernel_size=1, stride=1, padding=0, use_norm=SPECTRAL_IN_D)
 def forward(self, x):
x0 = x
x = self.conv1(x)
 x = self.activation(x)
 x = self.conv2(x)
 x = downsample(x)
 if RESIDUAL_IN_D:
x0 = downsample(x0)
 x0 = self.skip_conv(x0)
 x = x0 + x
return x
class Discriminator(nn.Module):
def __init__(self, num_classes: int = len(CLASSES), features: int = N_FEATURES, activation=DISCRIMINATOR_ACTIVATION)-> None:
super().__init__()
self.activation = activation
self.block0 = DiscriminatorDownOptimizedBlock(3, features * 1)
 if SELF_ATTENTION:
self.attn = SelfAttention(features * 1, use_norm=SPECTRAL_IN_D)
 self.block1 = DiscriminatorBlock(features * 1, features * 2)
 self.block2 = DiscriminatorBlock(features * 2, features * 4)
 self.block3 = DiscriminatorBlock(features * 4, features * 8)
 self.block4 = DiscriminatorBlock(features * 8, features * 8, down=False)
 self.linear = snlinear(features * 8, 1, use_norm=SPECTRAL_IN_D)
 if CONDITIONAL:
self.emb = sn_embedding(num_classes, features * 8, use_norm=SPECTRAL_IN_D)
 def forward(self, x, class_id):
x = self.block0(x)
 if SELF_ATTENTION:
x = self.attn(x)
 x = self.block1(x)
 x = self.block2(x)
 x = self.block3(x)
 x = self.block4(x)
 h = self.activation(x)
 h = h.sum(3 ).sum(2)
 output = self.linear(h ).squeeze(1)
 if CONDITIONAL:
h_labels = self.emb(class_id)
 output = output +(h * h_labels ).sum(1)
 return output
Discriminator()(torch.randn(4, 3, 64, 64), torch.zeros(4, dtype=torch.long)).shape<normalization>
dataset = DogDataset(data_path='.. /class_images/train/*.png',
transform=transforms.Compose([
transforms.RandomResizedCrop(64, scale=(0.7, 1.0), ratio=(1., 1.)),
transforms.RandomHorizontalFlip() ,
transforms.ToTensor() ,
transforms.Normalize(( 0.5, 0.5, 0.5),(0.5, 0.5, 0.5)) ,
]))
@torch.no_grad()
def inverse_transformation(tensor):
mean=[0.5, 0.5, 0.5]
std=[0.5, 0.5, 0.5]
for i in range(3):
m = mean[i]
s = std[i]
tensor[:, i].mul_(s ).add_(m)
 return tensor
@torch.no_grad()
def generate_images(generator, noises, labels=None):
if labels is None:
labels = torch.multinomial(torch.ones(len(CLASSES), dtype=torch.float32)/ len(CLASSES), noises.size(0)).to(noises.device)
 images = generator(noises, labels)
 return inverse_transformation(images)
 def show_images(tensor):
grid = torchvision.utils.make_grid(tensor)
 grid_numpy = grid.permute(1, 2, 0 ).cpu().numpy()
plt.imshow(grid_numpy)
 plt.show()<compute_test_metric>
dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size,
shuffle=True, num_workers=int(workers))
class DiscriminatorCriterion(nn.Module):
def forward(self, fake_images_output, real_images_output):
if LOSS == LossType.RaLS:
return torch.mean(( real_images_output - fake_images_output.mean() - 1)** 2)+ \
torch.mean(( fake_images_output - real_images_output.mean() + 1)** 2)
 elif LOSS == LossType.LS:
return torch.mean(real_images_output ** 2)+ torch.mean(( fake_images_output - 1)** 2)
 elif LOSS == LossType.Hinge:
return F.relu(1.0 - real_images_output ).mean() + F.relu(1 + fake_images_output ).mean()
else:
raise Exception("unimplemented: {}".format(LOSS))
class GeneratorCriterion(nn.Module):
def forward(self, fake_images_output, real_images_output):
if LOSS == LossType.RaLS:
return torch.mean(( real_images_output - fake_images_output.mean() + 1)** 2)+ \
torch.mean(( fake_images_output - real_images_output.mean() - 1)** 2)
 elif LOSS == LossType.LS:
return torch.mean(fake_images_output ** 2)
 elif LOSS == LossType.Hinge:
return -fake_images_output.mean()
else:
raise Exception("unimplemented: {}".format(LOSS))
def requires_grad(model, flag=True):
for p in model.parameters() :
p.requires_grad = flag
def warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor):
def f(x):
if x >= warmup_iters:
return 1
alpha = float(x)/ warmup_iters
return warmup_factor *(1 - alpha)+ alpha
return torch.optim.lr_scheduler.LambdaLR(optimizer, f)
 def train_one_epoch(epoch, generator, discriminator, generator_optimizer, discriminator_optimizer, dataloader, device, show_freq=None):
generator.train()
discriminator.train()
disc_criterion = DiscriminatorCriterion()
gen_criterion = GeneratorCriterion()
lr_scheduler_gen = None
lr_scheduler_disc = None
if False and epoch == 0:
lr_scheduler_gen = warmup_lr_scheduler(generator_optimizer, len(dataloader), warmup_factor=1./ 1000)
 lr_scheduler_disc = warmup_lr_scheduler(discriminator_optimizer, len(dataloader), warmup_factor=1./ 1000)
 loss_discriminators = []
loss_generators = []
for step,(real_images, real_labels)in enumerate(tqdm(dataloader, total=len(dataloader))):
discriminator.zero_grad()
real_images = real_images.to(device)
 real_labels = real_labels.to(device)
 batch_size = real_images.size(0)
 real_images_output = discriminator(real_images, real_labels)
 noise = sample(batch_size ).to(device)
 fake_images = generator(noise, real_labels)
 fake_images_output = discriminator(fake_images.detach() , real_labels)
 loss_discriminator = disc_criterion(fake_images_output, real_images_output)
 loss_discriminator.backward()
discriminator_optimizer.step()
loss_discriminators.append(loss_discriminator.detach().item())
 generator.zero_grad()
noise = sample(batch_size ).to(device)
 fake_images = generator(noise, real_labels)
 fake_images_output = discriminator(fake_images, real_labels)
 if LOSS.requires_real_for_generator() :
real_images_output = discriminator(real_images, real_labels ).detach()
else:
real_images_output = None
loss_generator = gen_criterion(fake_images_output, real_images_output)
 loss_generator.backward()
generator_optimizer.step()
loss_generators.append(loss_generator.detach().item())
 if show_freq and step % show_freq == 0:
generator.eval()
fake_images = generate_images(generator, FIXED_NOISES.to(device))
show_images(fake_images)
 generator.train()
if lr_scheduler_gen is not None:
lr_scheduler_gen.step()
if lr_scheduler_disc is not None:
lr_scheduler_disc.step()
print('[%d/%d] Loss_D: %.4f Loss_G: %.4f'
%(epoch + 1, EPOCHS,
np.mean(loss_discriminators), np.mean(loss_generators)) )<set_options>
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 print('device:', device )
gc.collect()
torch.cuda.empty_cache()
print(datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=9)) ).strftime("%Y/%m/%d %H:%M:%S"))
device = torch.device("cuda")if torch.cuda.is_available() else torch.device("cpu")
 generator = Generator().to(device)
 discriminator = Discriminator().to(device)
 generator_optimizer = torch.optim.Adam(generator.parameters() , lr=LR_GENERATOR, betas=(BETA1, BETA2))
discriminator_optimizer = torch.optim.Adam(discriminator.parameters() , lr=LR_DISCRIMINATOR, betas=(BETA1, BETA2))
if not os.path.exists("./tmp"):
os.mkdir("./tmp")
<train_model>
netG = Generator(nz=nz, nfeats=nch_g, nchannels=3 ).to(device)
 print(netG)
 netD = Discriminator(nchannels=3, nfeats=nch_d ).to(device)
 print(netD)
 criterion = nn.MSELoss()
optimizerD = optim.Adam(netD.parameters() , lr=lr, betas=(beta1, 0.999), weight_decay=1e-5)
 optimizerG = optim.Adam(netG.parameters() , lr=lr, betas=(beta1, 0.999), weight_decay=1e-5)
 fixed_noise = torch.randn(batch_size, nz, 1, 1, device=device)
 Loss_D_list, Loss_G_list = [], []
save_fake_image_count = 1
for epoch in range(EPOCHS):
print(f"[Epoch {epoch}]")
 train_one_epoch(epoch, generator, discriminator, generator_optimizer, discriminator_optimizer, dataloader, device, show_freq=SHOW_FREQ)
 generator.eval()
if(epoch + 1)% SHOW_FREQ_EPOCH == 0:
fake_images = generate_images(generator, FIXED_NOISES.to(device))
show_images(fake_images)
 if CONDITIONAL:
fake_images = generate_images(generator, FIXED_NOISES.to(device))
show_images(fake_images)
 save_image(fake_images, "./tmp/epoch{}.png".format(epoch + 1))<data_type_conversions>
def cosine_annealing(optimizer, start_lr, cur_steps, num_cycle):
t_cur = cur_steps % num_cycle
lr = 0.5 * start_lr *(math.cos(math.pi * t_cur / num_cycle)+ 1)
 for param_group in optimizer.param_groups:
param_group["lr"] = lr
return lr
print(datetime.datetime.now(tz=datetime.timezone(datetime.timedelta(hours=9)) ).strftime("%Y/%m/%d %H:%M:%S"))<import_modules>
def compute_gradient_penalty(D, X):

 lambda_gp = 10
alpha = torch.Tensor(np.random.random(size=X.shape)).cuda()
interpolates = alpha * X +(( 1 - alpha)*(X + 0.5 * X.std() * torch.rand(X.size() ).cuda()))
interpolates = Variable(interpolates, requires_grad=True ).cuda()
d_interpolates = D(interpolates)
 fake = Variable(torch.Tensor(X.shape[0], ).fill_(1.0), requires_grad=False ).cuda()
gradients = torch.autograd.grad(
 outputs=d_interpolates,
inputs=interpolates,
grad_outputs=fake,
create_graph=True,
retain_graph=True,
only_inputs=True, 
)[0]
gradient_penalty = lambda_gp *(( gradients.norm(2, dim=1)- 1)** 2 ).mean()
return gradient_penalty
import numpy as np
import pandas as pd
import os
import cv2
import random
import matplotlib.pyplot as plt
from PIL import Image
from skimage.util import random_noise
import xml.etree.ElementTree as ET
import torch
from torch import nn, optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torchvision.utils import save_image
from torch.autograd import Variable
from collections import OrderedDict
from torch.utils.data import TensorDataset, DataLoader,Dataset
from tqdm import tqdm_notebook as tqdm<set_options>
rap_time = start_time
for epoch in range(n_epoch):
for itr, data in tqdm(enumerate(dataloader), total = len(dataloader)) :
real_image = data.to(device)
 sample_size = real_image.size(0)
 noise = torch.randn(sample_size, nz, 1, 1, device=device)
 real_target = torch.full(( sample_size,), 1., device=device)
 fake_target = torch.full(( sample_size,), 0., device=device)
 netD.zero_grad()
output = netD(real_image)
 errD_real = criterion(output, real_target)
 D_x = output.mean().item()
fake_image = netG(noise)
 output = netD(fake_image.detach())
 errD_fake = criterion(output, fake_target)
 D_G_z1 = output.mean().item()
errD = errD_real + errD_fake
errD.backward()
optimizerD.step()
netG.zero_grad()
output = netD(fake_image)
 errG = criterion(output, real_target)
 errG.backward()
D_G_z2 = output.mean().item()
optimizerG.step()
writer.add_scalar("errD", errD.item() , epoch*len(dataloader)+itr)
 writer.add_scalar("errG", errG.item() , epoch*len(dataloader)+itr)
 writer.add_scalar("D_G_z1", D_G_z1, epoch*len(dataloader)+itr)
 writer.add_scalar("D_G_z2", D_G_z2, epoch*len(dataloader)+itr)
 writer.add_images('fake_images', fake_image[:10] * 0.5 + 0.5, global_step=epoch*len(dataloader)+itr, walltime=None, dataformats='NCHW')
 cosine_annealing(optimizerG, 0.0008, time.time() - start_time, 60*60*10)
 cosine_annealing(optimizerD, 0.0002, time.time() - start_time, 60*60*10)
 if time.time() - start_time > 60*60*8.95:
break
def seed_everything(seed=19960720):
random.seed(seed)
 os.environ['PYTHONHASHSEED'] = str(seed)
 np.random.seed(seed)
 torch.manual_seed(seed)
 torch.cuda.manual_seed(seed)
 torch.backends.cudnn.deterministic = True
seed_everything()<categorify>
def truncated_normal_(tensor, mean=0, std=1):
size = tensor.shape
tmp = tensor.new_empty(size +(4,)).normal_()
valid =(tmp < 2)&(tmp > -2)
 ind = valid.max(-1, keepdim=True)[1]
tensor.data.copy_(tmp.gather(-1, ind ).squeeze(-1))
tensor.data.mul_(std ).add_(mean)
 return tensor
def summary(model, input_size, batch_size=-1, device="cuda"):
def register_hook(module):
def hook(module, input, output):
class_name = str(module.__class__ ).split(".")[-1].split("'")[0]
module_idx = len(summary)
 m_key = "%s-%i" %(class_name, module_idx + 1)
 summary[m_key] = OrderedDict()
summary[m_key]["input_shape"] = list(input[0].size())
 summary[m_key]["input_shape"][0] = batch_size
if isinstance(output,(list, tuple)) :
summary[m_key]["output_shape"] = [
[-1] + list(o.size())[1:] for o in output
]
else:
summary[m_key]["output_shape"] = list(output.size())
 summary[m_key]["output_shape"][0] = batch_size
params = 0
if hasattr(module, "weight")and hasattr(module.weight, "size"):
params += torch.prod(torch.LongTensor(list(module.weight.size())))
 summary[m_key]["trainable"] = module.weight.requires_grad
if hasattr(module, "bias")and hasattr(module.bias, "size"):
params += torch.prod(torch.LongTensor(list(module.bias.size())))
 summary[m_key]["nb_params"] = params
if(
 not isinstance(module, nn.Sequential)
 and not isinstance(module, nn.ModuleList)
 and not(module == model)
):
hooks.append(module.register_forward_hook(hook))
device = device.lower()
assert device in [
"cuda",
"cpu",
], "Input device is not valid, please specify 'cuda' or 'cpu'"
if device == "cuda" and torch.cuda.is_available() :
dtype = torch.cuda.FloatTensor
else:
dtype = torch.FloatTensor
if isinstance(input_size, tuple):
input_size = [input_size]
x = [torch.rand(2, *in_size ).type(dtype)for in_size in input_size]
summary = OrderedDict()
hooks = []
model.apply(register_hook)
 model(*x)
 for h in hooks:
h.remove()
print("----------------------------------------------------------------")
 line_new = "{:>20} {:>25} {:>15}".format("Layer(type)", "Output Shape", "Param
print(line_new)
 print("================================================================")
 total_params = 0
total_output = 0
trainable_params = 0
for layer in summary:
line_new = "{:>20} {:>25} {:>15}".format(
 layer,
str(summary[layer]["output_shape"]),
"{0:,}".format(summary[layer]["nb_params"]), 
)
 total_params += summary[layer]["nb_params"]
total_output += np.prod(summary[layer]["output_shape"])
 if "trainable" in summary[layer]:
if summary[layer]["trainable"] == True:
trainable_params += summary[layer]["nb_params"]
print(line_new)
 total_input_size = abs(np.prod(input_size)* batch_size * 4./(1024 ** 2.))
 total_output_size = abs(2.* total_output * 4./(1024 ** 2.))
 total_params_size = abs(total_params.numpy() * 4./(1024 ** 2.))
 total_size = total_params_size + total_output_size + total_input_size
print("================================================================")
 print("Total params: {0:,}".format(total_params))
print("Trainable params: {0:,}".format(trainable_params))
print("Non-trainable params: {0:,}".format(total_params - trainable_params))
print("----------------------------------------------------------------")
 print("Input size(MB): %0.2f" % total_input_size)
 print("Forward/backward pass size(MB): %0.2f" % total_output_size)
 print("Params size(MB): %0.2f" % total_params_size)
 print("Estimated Total Size(MB): %0.2f" % total_size)
 print("----------------------------------------------------------------")
<normalization>
def truncated_normal(size, threshold=1):
values = truncnorm.rvs(-threshold, threshold, size=size)
 return torch.from_numpy(values ).float()
<categorify><EOS> 
torch.save(netG.state_dict() , 'netG.pth')
 torch.save(netD.state_dict() , 'netD.pth' )