resnet_id.py

import torch
import numpy as np
from matplotlib import pyplot as plt

torch.set_printoptions(precision=10)

class ResidualBlock(torch.nn.Module):
    def __init__(self, dims, bottleneck):
        super(ResidualBlock, self).__init__()
        self.linear1 = torch.nn.Linear(dims, bottleneck)
        self.linear2 = torch.nn.Linear(bottleneck, dims)
    
    def forward(self, x):
        y = self.linear1(x.relu())
        y = self.linear2(y.relu())
        return x + y

class ResNet(torch.nn.Module):
    def __init__(self, inout_dims, dims, bottleneck, depth):
        super(ResNet, self).__init__()
        self.linear_in = torch.nn.Linear(inout_dims, dims)
        self.residuals = torch.nn.Sequential(*[ResidualBlock(dims, bottleneck) for i in range(depth)])
        self.linear_out = torch.nn.Linear(dims, inout_dims)
    
    def forward(self, x):
        x = self.linear_in(x)
        x = self.residuals(x)
        x = self.linear_out(x)
        return x

dataset = torch.distributions.Normal(0, 1)

model = ResNet(64, 1024, 256, 1).cuda()
opt = torch.optim.Adam(model.parameters(), 3e-4)

for it in range(1000000):
    x = dataset.sample([1024, 64]).cuda()
    loss = (model(x) - x).square().sum(1).mean()
    print(it, loss.item())
    opt.zero_grad()
    loss.backward()
    opt.step()