cuda as CLI option

Author: Ipsedo

Diff for: iterative_machine_teaching/__init__.py

@@ -13,4 +13,10 @@
     SurrogateTeacher,
     Teacher,
 )
-from .train import TeachingType, train
+from .train import (
+    DatasetOptions,
+    StudentOptions,
+    TeacherOptions,
+    TeachingType,
+    train,
+)

Diff for: iterative_machine_teaching/__main__.py

@@ -2,22 +2,68 @@
 import argparse
 
 from .data import load_gaussian, load_mnist
-from .train import TeachingType, train
+from .train import (
+    DatasetOptions,
+    StudentOptions,
+    TeacherOptions,
+    TeachingType,
+    train,
+)
 
 
 def main() -> None:
     parser = argparse.ArgumentParser("IterativeMachineTeaching")
 
+    parser.add_argument("--cuda", type=bool, action="store_true")
     parser.add_argument("kind", type=TeachingType, choices=list(TeachingType))
+    parser.add_argument("--train-ratio", type=float, default=4.0 / 5.0)
 
+    # student / example options
     parser.add_argument(
-        "-l",
-        "--limit-train",
+        "--student-examples",
         type=int,
         default=-1,
         help="Number of examples in student train dataset, "
         "negative value means max",
     )
+    parser.add_argument(
+        "--student-steps",
+        type=int,
+        default=1024,
+        help="Number of forward / backward steps for student",
+    )
+    parser.add_argument(
+        "--student-batch-size",
+        type=int,
+        default=8,
+        help="Batch size for student and example",
+    )
+    parser.add_argument(
+        "--student-lr",
+        type=float,
+        default=1e-3,
+        help="Student and example learning rate",
+    )
+
+    # teacher options
+    parser.add_argument(
+        "--teacher-lr", type=float, default=1e-3, help="Teacher learning rate"
+    )
+    parser.add_argument(
+        "--teacher-batch-size", type=int, default=8, help="Teacher batch size"
+    )
+    parser.add_argument(
+        "--research-batch-size",
+        type=int,
+        default=512,
+        help="Batch size for example research",
+    )
+    parser.add_argument(
+        "--teacher-epochs",
+        type=int,
+        default=16,
+        help="Teacher training epochs",
+    )
 
     dataset_subparser = parser.add_subparsers(
         title="dataset", dest="dataset", required=True
@@ -35,14 +81,36 @@ def main() -> None:
     args = parser.parse_args()
 
     if args.dataset == "mnist":
-        dataset = load_mnist(args.input_pickle)
+        x, y = load_mnist(args.input_pickle)
     elif args.dataset == "gaussian":
-        dataset = load_gaussian(args.dim, args.per_class_example)
+        x, y = load_gaussian(args.dim, args.per_class_example)
     else:
         parser.error("Unrecognized dataset")
         return
 
-    train(dataset, args.dataset, args.kind, args.limit_train)
+    dataset_options = DatasetOptions(args.dataset, x, y, args.train_ratio)
+
+    student_options = StudentOptions(
+        args.student_examples,
+        args.student_steps,
+        args.student_batch_size,
+        args.student_lr,
+    )
+
+    teacher_options = TeacherOptions(
+        args.teacher_lr,
+        args.teacher_batch_size,
+        args.research_batch_size,
+        args.teacher_epochs,
+    )
+
+    train(
+        dataset_options,
+        args.kind,
+        teacher_options,
+        student_options,
+        args.cuda,
+    )
 
 
 if __name__ == "__main__":

Diff for: iterative_machine_teaching/train.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 from enum import Enum
-from typing import Dict, Tuple, Type
+from typing import Dict, NamedTuple, Tuple, Type
 
 import matplotlib.pyplot as plt
 import torch as th
@@ -21,6 +21,36 @@
     Teacher,
 )
 
+DatasetOptions = NamedTuple(
+    "DatasetOptions",
+    [
+        ("name", str),
+        ("x", th.Tensor),
+        ("y", th.Tensor),
+        ("train_ratio", float),
+    ],
+)
+
+StudentOptions = NamedTuple(
+    "StudentOptions",
+    [
+        ("examples", int),
+        ("steps", int),
+        ("batch_size", int),
+        ("learning_rate", float),
+    ],
+)
+
+TeacherOptions = NamedTuple(
+    "TeacherOptions",
+    [
+        ("learning_rate", float),
+        ("batch_size", int),
+        ("research_batch_size", int),
+        ("nb_epoch", int),
+    ],
+)
+
 
 class TeachingType(Enum):
     OMNISCIENT = "OMNISCIENT"
@@ -49,56 +79,68 @@ def get_student(self, clf: Classifier, learning_rate: float) -> Student:
 
 
 def train(
-    dataset: Tuple[th.Tensor, th.Tensor],
-    dataset_name: str,
+    dataset_options: DatasetOptions,
     kind: TeachingType,
-    example_nb_student: int,
+    teacher_options: TeacherOptions,
+    student_options: StudentOptions,
+    cuda: bool,
 ) -> None:
 
-    x, y = dataset
+    assert 0.0 < dataset_options.train_ratio < 1.0
 
-    num_features = x.size()[1]  # 784
-    num_classes = th.unique(y).size()[0]  # 10
+    x, y = dataset_options.x, dataset_options.y
+
+    num_features = x.size()[1]
+    num_classes = th.unique(y).size()[0]
 
     print(
-        f'Dataset "{dataset_name}" of {x.size()[0]} '
+        f'Dataset "{dataset_options.name}" of {x.size()[0]} '
         f"examples with {kind.value} teacher."
     )
 
-    ratio_train = 4.0 / 5.0
-    limit_train = int(x.size()[0] * ratio_train)
+    limit_train = int(x.size()[0] * dataset_options.train_ratio)
 
-    x_train = x[:limit_train, :].cuda()
-    y_train = y[:limit_train].cuda()
+    x_train = x[:limit_train, :]
+    y_train = y[:limit_train]
 
-    x_test = x[limit_train:, :].cuda()
-    y_test = y[limit_train:].cuda()
+    x_test = x[limit_train:, :]
+    y_test = y[limit_train:]
 
     # create models
-    student_model = LinearClassifier(num_features, num_classes).cuda()
-    teacher_model = LinearClassifier(num_features, num_classes).cuda()
+    student_model = LinearClassifier(num_features, num_classes)
+    example_model = LinearClassifier(num_features, num_classes)
+    teacher_model = LinearClassifier(num_features, num_classes)
+
+    # cuda or not
+    if cuda:
+        x_train = x_train.cuda()
+        y_train = y_train.cuda()
+
+        x_test = x_test.cuda()
+        y_test = y_test.cuda()
 
-    # create student and teacher
-    learning_rate = 1e-3
-    research_batch_size = 512
+        student_model = student_model.cuda()
+        example_model = example_model.cuda()
+        teacher_model = teacher_model.cuda()
 
-    student = kind.get_student(student_model, learning_rate)
+    # create student, example and teacher
+    student = kind.get_student(student_model, student_options.learning_rate)
+    example = ModelWrapper(example_model, student_options.learning_rate)
     teacher = kind.get_teacher(
-        teacher_model, learning_rate, research_batch_size
+        teacher_model,
+        teacher_options.learning_rate,
+        teacher_options.research_batch_size,
     )
 
     # Train teacher
     print("Train teacher...")
+    nb_batch_teacher = x_train.size()[0] // teacher_options.batch_size
 
-    nb_epoch_teacher = 25
-    batch_size_teacher = 32
-    nb_batch_teacher = x_train.size()[0] // batch_size_teacher
-
-    tqdm_bar = tqdm(range(nb_epoch_teacher))
+    tqdm_bar = tqdm(range(teacher_options.nb_epoch))
     for e in tqdm_bar:
         for b_idx in range(nb_batch_teacher):
-            i_min = b_idx * batch_size_teacher
-            i_max = (b_idx + 1) * batch_size_teacher
+            i_min = b_idx * teacher_options.batch_size
+            i_max = (b_idx + 1) * teacher_options.batch_size
 
             _ = teacher.train(x_train[i_min:i_max], y_train[i_min:i_max])
 
@@ -109,35 +151,31 @@ def train(
 
         tqdm_bar.set_description(f"Epoch {e} : F1-Score = {f1_score_value}")
 
-    # For comparison
+    # For benchmark
 
     # to avoid a lot of compute...
     # if negative -> all train examples
-    example_nb_student = (
-        example_nb_student if example_nb_student >= 0 else x_train.size()[0]
+    student_examples = (
+        student_options.examples
+        if student_options.examples >= 0
+        else x_train.size()[0]
     )
-    x_train = x_train[:example_nb_student]
-    y_train = y_train[:example_nb_student]
+    x_train = x_train[:student_examples]
+    y_train = y_train[:student_examples]
 
-    rounds = 1024
-    batch_size = 16
-    nb_batch = x_train.size()[0] // batch_size
+    nb_batch = x_train.size()[0] // student_options.batch_size
 
     # train example
     print("Train example...")
 
-    example = ModelWrapper(
-        LinearClassifier(num_features, num_classes).cuda(), learning_rate
-    )
-
     batch_index_example = 0
     loss_values_example = []
     metrics_example = []
 
-    for _ in tqdm(range(rounds)):
+    for _ in tqdm(range(student_options.steps)):
         b_idx = batch_index_example % nb_batch
-        i_min = b_idx * batch_size
-        i_max = (b_idx + 1) * batch_size
+        i_min = b_idx * student_options.batch_size
+        i_max = (b_idx + 1) * student_options.batch_size
 
         loss = example.train(x_train[i_min:i_max], y_train[i_min:i_max])
 
@@ -158,9 +196,9 @@ def train(
     loss_values_student = []
     metrics_student = []
 
-    for _ in tqdm(range(rounds)):
+    for _ in tqdm(range(student_options.steps)):
         selected_x, selected_y = teacher.select_n_examples(
-            student, x_train, y_train, batch_size
+            student, x_train, y_train, student_options.batch_size
         )
 
         loss = student.train(selected_x, selected_y)
@@ -180,7 +218,7 @@ def train(
     plt.plot(metrics_example, c="blue", label="example - f1 score")
     plt.plot(metrics_student, c="red", label="student - f1 score")
 
-    plt.title(f"{dataset_name} Linear - {kind.value}")
+    plt.title(f"{dataset_options.name} Linear - {kind.value}")
     plt.xlabel("mini-batch optim steps")
     plt.legend()
     plt.show()

Diff for: setup.py

@@ -3,12 +3,10 @@
 
 setup(
     name="iterative_machine_teaching",
-    version="1.0.0",
+    version="1.1.0",
     packages=["iterative_machine_teaching"],
     url="https://github.com/Ipsedo/IterativeMachineTeaching",
-    license="",
+    license="GPL-3.0",
     author="Ipsedo",
-    author_email="",
     description="Iterative Machine Teaching implementation",
-    test_suite="tests",
 )