数据、破坏与瓶颈 · 实现代码

def digit_data():
    data=load_digits(); x=torch.tensor(data.images/16,dtype=torch.float32)[:,None]
    train,test=train_test_split(np.arange(len(x)),test_size=.3,random_state=SEED,stratify=data.target)
    assert set(train).isdisjoint(test)
    return x,data.target,train,test

def noise_corrupt(x):
    return (x*(torch.rand_like(x)>.25)+torch.randn_like(x)*.15).clamp(0,1)

def patch_mask(n,ratio=.5):
    # Grid is 4x4 patches, each 2x2 pixels. 1 means HIDDEN.
    order=torch.rand(n,16).argsort(dim=1)
    keep=int(16*(1-ratio))
    patch=torch.ones(n,16);patch.scatter_(1,order[:,:keep],0)
    return patch.reshape(n,1,4,4).repeat_interleave(2,2).repeat_interleave(2,3)

class ConvAE(nn.Module):
    def __init__(self,channels=1):
        super().__init__()
        self.encoder=nn.Sequential(nn.Conv2d(channels,8,3,stride=2,padding=1),nn.ReLU(),
          nn.Conv2d(8,16,3,stride=2,padding=1),nn.ReLU(),nn.Flatten(),nn.Linear(64,16))
        self.expand=nn.Linear(16,64)
        self.decoder=nn.Sequential(nn.ConvTranspose2d(16,8,4,stride=2,padding=1),nn.ReLU(),
          nn.ConvTranspose2d(8,1,4,stride=2,padding=1),nn.Sigmoid())
    def forward(self,x):
        h=self.encoder(x);return self.decoder(self.expand(h).reshape(-1,16,2,2)),h