def train_model(train_data, train_labels, cls_names, dataset, verbose=False):
path_manager = PathManager(MLPNet.NAME, dataset)
logger.info("Train the model: {}".format("-".join(cls_names)))
train_loader = list(
gen_batch(train_data, train_labels, batch_size=BATCH_SIZE))
model = MLPNet(len(cls_names),
train_data[0, :].size()[0]).to(config.DEVICE)
optimizer = optim.Adam(model.parameters())
# optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
loss_func = nn.CrossEntropyLoss()
for epoch in trange(1, EPOCH + 1):
total_loss = 0
for batch_idx in trange(0, len(train_loader)):
x_tensor, target_tensor = train_loader[batch_idx]
optimizer.zero_grad()
loss = loss_func(model.logit(x_tensor), target_tensor)
total_loss = total_loss + loss.data.item()
loss.backward()
optimizer.step()
if (batch_idx + 1) % 5 == 0 and verbose:
logger.info(
'=> epoch: {}, batch index: {}, train loss: {:.6f}'.format(
epoch, batch_idx + 1, total_loss))
logger.info("Save model path => {}".format(path_manager.mdl_path()))
torch.save(model.state_dict(), path_manager.mdl_path())
logger.info("Finish training")
return model
def evaluate_mnist():
from openapi.utils import MNIST
IMG_SIZE = 28
cls_names = MNIST._ID2STR
path_manager = PathManager(LinearModelTree.NAME, MNIST.NAME)
fmnist = MNIST(cls_names)
_, _, test_imgs, test_labels = fmnist.load_mnist()
test_imgs = test_imgs.view(test_imgs.size()[0], IMG_SIZE * IMG_SIZE)
mdl = load_model(path_manager.mdl_path())
evaluate(mdl, test_imgs, test_labels, "test")
def train_model(train_data, train_labels, cls_names, dataset, verbose=False):
logger.info("Train the model: {}".format("-".join(cls_names)))
path_manager = PathManager(LinearModelTree.NAME, dataset)
var_num = train_data[0, :].size()[0]
lmt = LinearModelTree(min_node_size=100,
clss_num=len(cls_names),
var_num=var_num)
lmt.build_tree(train_data, train_labels)
logger.info("Finish training Node {}".format(lmt.node))
with open(path_manager.mdl_path(), "wb") as f:
pickle.dump(lmt, f)
return lmt
def evaluate_fmnist():
from openapi.utils import FMNIST
IMG_SIZE = 28
# cls_names = ["Shirt", "Pullover", "Coat"]
cls_names = FMNIST._ID2STR
path_manager = PathManager(MLPNet.NAME, FMNIST.NAME)
fmnist = FMNIST(cls_names)
_, _, test_imgs, test_labels = fmnist.load_mnist()
test_imgs = test_imgs.view(test_imgs.size()[0], IMG_SIZE * IMG_SIZE)
mdl = load_model(path_manager.mdl_path(), len(cls_names),
IMG_SIZE * IMG_SIZE)
evaluate(mdl, test_imgs, test_labels, "test")
def evaluate_fmnist():
IMG_SIZE = 28
cls_names = [
"T-shirt_top", "Trouser", "Pullover", "Dress", "Coat", "Sandal",
"Shirt", "Sneaker", "Bag", "Ankle_boot"
]
path_manager = PathManager(NAME, "FMNIST")
fmnist = FMNIST(cls_names)
train_imgs, train_labels, test_imgs, test_labels = fmnist.load_mnist()
fmnist.display(train_labels, test_labels)
test_imgs = test_imgs.view(test_imgs.size()[0], IMG_SIZE * IMG_SIZE)
mdl = load_model(path_manager.mdl_path(), len(cls_names),
IMG_SIZE * IMG_SIZE)
evaluate(mdl, test_imgs, test_labels, "test")
def train_mnist():
from openapi.utils import MNIST
IMG_SIZE = 28
cls_names = MNIST._ID2STR
mnist = MNIST(cls_names)
path_manager = PathManager(MLPNet.NAME, MNIST.NAME)
train_imgs, train_labels, test_imgs, test_labels = mnist.load_mnist()
mnist.display(train_labels, test_labels)
train_imgs = train_imgs.view(train_imgs.size()[0], IMG_SIZE * IMG_SIZE)
train_model(train_imgs,
train_labels,
cls_names,
MNIST.NAME,
verbose=True)
mdl = load_model(path_manager.mdl_path(), len(cls_names),
IMG_SIZE * IMG_SIZE)
evaluate(mdl, train_imgs, train_labels, "train")
def train_fmnist():
from openapi.utils import FMNIST
IMG_SIZE = 28
# cls_names = ["Shirt", "Pullover", "Coat"]
cls_names = FMNIST._ID2STR
fmnist = FMNIST(cls_names)
path_manager = PathManager(MLPNet.NAME, FMNIST.NAME)
train_imgs, train_labels, test_imgs, test_labels = fmnist.load_mnist()
fmnist.display(train_labels, test_labels)
train_imgs = train_imgs.view(train_imgs.size()[0], IMG_SIZE * IMG_SIZE)
train_model(train_imgs,
train_labels,
cls_names,
"FMNIST",
verbose=True)
mdl = load_model(path_manager.mdl_path(), len(cls_names),
IMG_SIZE * IMG_SIZE)
evaluate(mdl, train_imgs, train_labels, "train")