def run(num_epochs=5, device="cuda"):
criterion = nn.CrossEntropyLoss()
augmenter = Augmenter(approximate=True)
full_dataset = DBDataset(TrainImage, MathSymbol, get_data(augmenter),
get_label, get_class_name)
full_dataset = BalancedDS(full_dataset)
test_train_split = 0.9
train_size = int(test_train_split * len(full_dataset))
test_size = len(full_dataset) - train_size
train_dataset, test_dataset = random_split(full_dataset,
[train_size, test_size])
dataloaders = {
"train": DataLoader(train_dataset, batch_size=32, num_workers=1),
"test": DataLoader(test_dataset, batch_size=16, shuffle=True,
num_workers=1)
}
dataset_sizes = {
"train": len(train_dataset),
"test": len(test_dataset)
}
print(dataloaders)
print(dataset_sizes)
model = MobileNet(features=len(full_dataset.classes), pretrained=True)
model.freeze()
model = model.to(device)
def unfreeze(x, y):
model.unfreeze()
augmenter.approximate = False
solver = Solver(criterion, dataloaders, model, cb=unfreeze)
model, accuracy = solver.train(device=device,
num_epochs=num_epochs)
eval_model(model, dataloaders["test"], device, len(full_dataset.classes))
model = model.to('cpu')
torch.save(model.state_dict(), "test_augment.pth")
model.estimate_variane = True
byteArr = model.to_onnx()
torchByteArr = io.BytesIO()
torch.save(model.state_dict(), torchByteArr)
model_entity = ClassificationModel(None, model=byteArr.getvalue(),
timestamp=timezone.now(),
pytorch=torchByteArr.getvalue(),
accuracy=accuracy)
model_entity.save()
def test_can_export_and_load_mobilenet(self):
model = MobileNet(features=20)
f = 'model.pth'
model.to_file(f)
model_2 = MobileNet.from_file(f)
self.assertEquals(model.num_features, model_2.num_features)
def to_pytorch(self) -> MobileNet:
state_dict = torch.load(io.BytesIO(self.pytorch),
map_location=torch.device('cpu'))
model = MobileNet(
features=state_dict['mobilenet.classifier.1.bias'].shape[0],
pretrained=False)
model.load_state_dict(state_dict)
return model
def test_mobilenet_can_use_testtime_dropout(self):
model = MobileNet(features=20, test_time_dropout=True)
model = model.eval()
inp = torch.randn(1, 3, 224, 224, device="cpu")
result1 = model(inp)
result2 = model(inp)
self.assertFalse(torch.all(result1.eq(result2)))
def run():
class_table = MathSymbol
classes = [
get_class_name(cls_ent)
for cls_ent in class_table.objects.all().order_by('timestamp')
]
class_to_ix = {
get_class_name(cls_ent): ix
for ix, cls_ent in enumerate(class_table.objects.all().order_by(
'timestamp'))
}
print(class_to_ix)
file_name = "test_augment.pth"
model = MobileNet.from_file(file_name,
test_time_dropout=False,
estimate_variane=True)
model.eval().cuda()
iters = 200
full_dataset = datasets.ImageFolder("res/certain", preprocess)
dataloader = DataLoader(full_dataset,
batch_size=1,
shuffle=False,
num_workers=4)
for i, data in enumerate(dataloader):
preds = torch.zeros((iters, 63))
softmaxs = torch.zeros((iters, 63))
maxs = torch.zeros(iters)
variances = None
inputs, labels = data
print(labels)
inputs = inputs.cuda()
for j in range(iters):
res = model(inputs)
pred = res[0]
var = res[1]
pred, var = pred.detach().cpu(), var.detach().cpu()
preds[j] = pred[0]
maxs[j] = pred.argmax()
softmaxs[j] = F.softmax(pred[0])
variances = var[0]
stds = softmaxs.std(dim=0)
m = int(maxs[0].item())
print(m)
print(stds[m] / 0.001)
print(variances[m] / 0.001)
def get_upload_json(file_name, **kwargs):
pytorch = Path(file_name).read_bytes()
model = MobileNet.from_file(file_name, **kwargs)
model.eval()
byte_arr = model.to_onnx()
json = {
"pytorch": base64.b64encode(pytorch).decode('utf-8'),
"onnx": base64.b64encode(byte_arr.getvalue()).decode('utf-8')
}
return json
def run():
path = 'test_augment.pth'
pytorch = Path(path).read_bytes()
model = MobileNet.from_file(path)
model.eval()
byte_arr = model.to_onnx()
model_instance = ClassificationModel(None,
model=byte_arr.getvalue(),
pytorch=pytorch,
timestamp=timezone.now(),
accuracy=0.99)
model_instance.save()
def run():
file_name = "test_augment.pth"
model = MobileNet.from_file(file_name,
test_time_dropout=False,
estimate_variane=True)
byteArr = model.to_onnx()
torchByteArr = io.BytesIO()
torch.save(model.state_dict(), torchByteArr)
model_entity = ClassificationModel(None,
model=byteArr.getvalue(),
timestamp=timezone.now(),
pytorch=torchByteArr.getvalue(),
accuracy=0.99)
model_entity.save()
def update_train_features(torch_model, num_classes):
with torch.no_grad():
model = MobileNet(features=num_classes, pretrained=False)
model.load_state_dict(
torch.load(torch_model, map_location=torch.device('cpu')))
model = model.eval()
for i, train_img in enumerate(TrainImage.objects.all()):
if i % 10 == 0:
print(f'Updating image {i}')
image = Image.open(io.BytesIO(train_img.image))
data = preprocess(image)
img = data.repeat((3, 1, 1))
img = img.reshape((1, img.shape[0], img.shape[1], img.shape[2]))
features = model.features(img)
features = features.mean([2, 3])
byte_f = io.BytesIO()
torch.save(features, byte_f)
train_img.features = byte_f.getvalue()
train_img.save()
def test_mobilenet_can_predict_uncertainty(self):
model = MobileNet(features=20, estimate_variane=True)
inp = torch.randn(1, 3, 224, 224, device="cpu")
result = model(inp)
self.assertEquals(list(result.shape), [2, 1, 20])