Diffusion unet

import os
os.environ['CUDA_VISIBLE_DEVICES']='1'

torch.set_printoptions(precision=4, linewidth=140, sci_mode=False)
torch.manual_seed(1)
mpl.rcParams['image.cmap'] = 'gray_r'
mpl.rcParams['figure.dpi'] = 70

import logging
logging.disable(logging.WARNING)

set_seed(42)
if fc.defaults.cpus>8: fc.defaults.cpus=8

xl,yl = 'image','label'
name = "fashion_mnist"
bs = 512
dsd = load_dataset(name)

@inplace
def transformi(b): b[xl] = [F.pad(TF.to_tensor(o), (2,2,2,2))-0.5 for o in b[xl]]

tds = dsd.with_transform(transformi)
dls = DataLoaders(dl_ddpm(tds['train']), dl_ddpm(tds['test']))

dl = dls.train
(xt,t),eps = b = next(iter(dl))

Train

Based on Diffusers

# This version is giving poor results - use the cell below instead
class SelfAttention(nn.Module):
    def __init__(self, ni, attn_chans):
        super().__init__()
        self.attn = nn.MultiheadAttention(ni, ni//attn_chans, batch_first=True)
        self.norm = nn.BatchNorm2d(ni)

    def forward(self, x):
        n,c,h,w = x.shape
        x = self.norm(x).view(n, c, -1).transpose(1, 2)
        x = self.attn(x, x, x, need_weights=False)[0]
        return x.transpose(1,2).reshape(n,c,h,w)

lr = 1e-2
epochs = 25
opt_func = partial(optim.Adam, eps=1e-5)
tmax = epochs * len(dls.train)
sched = partial(lr_scheduler.OneCycleLR, max_lr=lr, total_steps=tmax)
cbs = [DeviceCB(), ProgressCB(plot=True), MetricsCB(), BatchSchedCB(sched), MixedPrecision()]
model = EmbUNetModel(in_channels=1, out_channels=1, nfs=(32,64,128,256), num_layers=2)
learn = Learner(model, dls, nn.MSELoss(), lr=lr, cbs=cbs, opt_func=opt_func)

learn.fit(epochs)

loss	epoch	train
0.150	0	train
0.086	0	eval
0.069	1	train
0.171	1	eval
0.057	2	train
0.071	2	eval
0.050	3	train
0.055	3	eval
0.045	4	train
0.050	4	eval
0.043	5	train
0.073	5	eval
0.041	6	train
0.044	6	eval
0.039	7	train
0.044	7	eval
0.038	8	train
0.043	8	eval
0.038	9	train
0.058	9	eval
0.038	10	train
0.044	10	eval
0.036	11	train
0.042	11	eval
0.035	12	train
0.038	12	eval
0.035	13	train
0.039	13	eval
0.034	14	train
0.036	14	eval
0.034	15	train
0.036	15	eval
0.034	16	train
0.034	16	eval
0.034	17	train
0.035	17	eval
0.033	18	train
0.033	18	eval
0.033	19	train
0.033	19	eval
0.033	20	train
0.033	20	eval
0.033	21	train
0.032	21	eval
0.032	22	train
0.034	22	eval
0.032	23	train
0.032	23	eval
0.032	24	train
0.033	24	eval

Sampling

from miniai.fid import ImageEval

cmodel = torch.load('models/data_aug2.pkl')
del(cmodel[8])
del(cmodel[7])

@inplace
def transformi(b): b[xl] = [F.pad(TF.to_tensor(o), (2,2,2,2))*2-1 for o in b[xl]]

bs = 2048
tds = dsd.with_transform(transformi)
dls = DataLoaders.from_dd(tds, bs, num_workers=fc.defaults.cpus)

dt = dls.train
xb,yb = next(iter(dt))

ie = ImageEval(cmodel, dls, cbs=[DeviceCB()])

sz = (2048,1,32,32)

# set_seed(42)
preds = sample(ddim_step, model, sz, steps=100, eta=1.)
s = (preds[-1]*2)
s.min(),s.max(),s.shape

100.00% [100/100 00:53<00:00]

(tensor(-1.0918), tensor(1.4292), torch.Size([2048, 1, 32, 32]))

show_images(s[:25].clamp(-1,1), imsize=1.5)

ie.fid(s),ie.kid(s),s.shape

(4.058064770194278, 0.010895456187427044, torch.Size([2048, 1, 32, 32]))

preds = sample(ddim_step, model, sz, steps=100, eta=1.)
ie.fid(preds[-1]*2)

100.00% [100/100 00:53<00:00]

5.320260029850715

preds = sample(ddim_step, model, sz, steps=50, eta=1.)
ie.fid(preds[-1]*2)

100.00% [50/50 00:26<00:00]

5.243807277315682

preds = sample(ddim_step, model, sz, steps=50, eta=1.)
ie.fid(preds[-1]*2)

100.00% [50/50 00:26<00:00]

4.963977301033992

Conditional model

def collate_ddpm(b):
    b = default_collate(b)
    (xt,t),eps = noisify(b[xl])
    return (xt,t,b[yl]),eps

@inplace
def transformi(b): b[xl] = [F.pad(TF.to_tensor(o), (2,2,2,2))-0.5 for o in b[xl]]

tds = dsd.with_transform(transformi)
dls = DataLoaders(dl_ddpm(tds['train']), dl_ddpm(tds['test']))

dl = dls.train
(xt,t,c),eps = b = next(iter(dl))

class CondUNetModel(nn.Module):
    def __init__( self, n_classes, in_channels=3, out_channels=3, nfs=(224,448,672,896), num_layers=1):
        super().__init__()
        self.conv_in = nn.Conv2d(in_channels, nfs[0], kernel_size=3, padding=1)
        self.n_temb = nf = nfs[0]
        n_emb = nf*4
        self.cond_emb = nn.Embedding(n_classes, n_emb)
        self.emb_mlp = nn.Sequential(lin(self.n_temb, n_emb, norm=nn.BatchNorm1d),
                                     lin(n_emb, n_emb))
        self.downs = nn.ModuleList()
        for i in range(len(nfs)):
            ni = nf
            nf = nfs[i]
            self.downs.append(DownBlock(n_emb, ni, nf, add_down=i!=len(nfs)-1, num_layers=num_layers))
        self.mid_block = EmbResBlock(n_emb, nfs[-1])

        rev_nfs = list(reversed(nfs))
        nf = rev_nfs[0]
        self.ups = nn.ModuleList()
        for i in range(len(nfs)):
            prev_nf = nf
            nf = rev_nfs[i]
            ni = rev_nfs[min(i+1, len(nfs)-1)]
            self.ups.append(UpBlock(n_emb, ni, prev_nf, nf, add_up=i!=len(nfs)-1, num_layers=num_layers+1))
        self.conv_out = pre_conv(nfs[0], out_channels, act=nn.SiLU, norm=nn.BatchNorm2d, bias=False)

    def forward(self, inp):
        x,t,c = inp
        temb = timestep_embedding(t, self.n_temb)
        cemb = self.cond_emb(c)
        emb = self.emb_mlp(temb) + cemb
        x = self.conv_in(x)
        saved = [x]
        for block in self.downs: x = block(x, emb)
        saved += [p for o in self.downs for p in o.saved]
        x = self.mid_block(x, emb)
        for block in self.ups: x = block(x, emb, saved)
        return self.conv_out(x)

lr = 1e-2
epochs = 25
opt_func = partial(optim.Adam, eps=1e-5)
tmax = epochs * len(dls.train)
sched = partial(lr_scheduler.OneCycleLR, max_lr=lr, total_steps=tmax)
cbs = [DeviceCB(), ProgressCB(plot=True), MetricsCB(), BatchSchedCB(sched), MixedPrecision()]
model = CondUNetModel(10, in_channels=1, out_channels=1, nfs=(32,64,128,256), num_layers=2)
learn = Learner(model, dls, nn.MSELoss(), lr=lr, cbs=cbs, opt_func=opt_func)

learn.fit(epochs)

loss	epoch	train
0.178	0	train
0.099	0	eval
0.072	1	train
0.066	1	eval
0.053	2	train
0.053	2	eval
0.047	3	train
0.050	3	eval
0.045	4	train
0.045	4	eval
0.042	5	train
0.048	5	eval
0.041	6	train
0.060	6	eval
0.039	7	train
0.042	7	eval
0.037	8	train
0.039	8	eval
0.037	9	train
0.051	9	eval
0.036	10	train
0.039	10	eval
0.035	11	train
0.041	11	eval
0.035	12	train
0.041	12	eval
0.034	13	train
0.035	13	eval
0.034	14	train
0.035	14	eval
0.034	15	train
0.036	15	eval
0.033	16	train
0.037	16	eval
0.033	17	train
0.032	17	eval
0.032	18	train
0.036	18	eval
0.032	19	train
0.033	19	eval
0.032	20	train
0.032	20	eval
0.032	21	train
0.033	21	eval
0.032	22	train
0.033	22	eval
0.031	23	train
0.032	23	eval
0.031	24	train
0.033	24	eval

sz = (256,1,32,32)

lbls = dsd['train'].features[yl].names
lbls

['T - shirt / top',
 'Trouser',
 'Pullover',
 'Dress',
 'Coat',
 'Sandal',
 'Shirt',
 'Sneaker',
 'Bag',
 'Ankle boot']

set_seed(42)
cid = 0
preds = sample(cid, ddim_step, model, sz, steps=100, eta=1.)
s = (preds[-1]*2)
show_images(s[:25].clamp(-1,1), imsize=1.5, suptitle=lbls[cid])

100.00% [100/100 00:02<00:00]

set_seed(42)
cid = 0
preds = sample(cid, ddim_step, model, sz, steps=100, eta=0.)
s = (preds[-1]*2)
show_images(s[:25].clamp(-1,1), imsize=1.5, suptitle=lbls[cid])

100.00% [100/100 00:02<00:00]