class CNN(nn.Module):
def __init__(self, config):
super(CNN, self).__init__()
self.module = Learner(get_cnn(config))
def forward(self, x, vars=None, bn_training=True):
if vars is None:
return self.module(x)
else:
return self.module.functional(vars, bn_training, x)
def test_resnet():
try:
from torchvision.models.resnet import resnet18
from metann import Learner
net = resnet18()
net = Learner(net)
print(net.functional(net.parameters(), True, torch.randn(3, 3, 224, 224)))
except ImportError:
Warning('torchvision not included, cannot be tested')
return
finally:
return
def test_learner():
net = Learner(
nn.Sequential(
nn.Conv2d(3, 3, 3),
nn.Conv2d(3, 3, 3),
Flatten(),
nn.Linear(3, 4),
)).to(device)
x = torch.randn(3, 3, 5, 5).to(device)
y = torch.randint(0, 4, (3, )).to(device)
criterion = nn.CrossEntropyLoss()
params = list(net.parameters())
for i in range(500):
outs = net.functional(params, True, x)
loss = criterion(outs, y)
grads = torch.autograd.grad(loss, params)
with torch.no_grad():
params = [(a - 0.01 * b).requires_grad_()
for a, b in zip(params, grads)]
print(loss)
assert loss <= 0.05