def main(): # Training settings parser = argparse.ArgumentParser() parser.add_argument('--batch-size', type=int, default=64, metavar='N', help='input batch size for training (default: 64)') parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N', help='input batch size for testing (default: 1000)') parser.add_argument('--epochs', type=int, default=10, metavar='N', help='number of epochs to train (default: 10)') parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)') parser.add_argument('--momentum', type=float, default=0.5, metavar='M', help='SGD momentum (default: 0.5)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument( '--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status') args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") train_loader, test_loader = get_data_loader(args.batch_size, args.test_batch_size, use_cuda) model = load_model(device=device) optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) # training for epoch in range(1, args.epochs + 1): train(model, device, train_loader, optimizer, epoch, args.log_interval) test(model, device, test_loader) # model serializing model_dir = PathManager.MODEL_DIR model_dir.mkdir(parents=True, exist_ok=True) model_path = model_dir / "mnist_cnn.pth" torch.save(model.state_dict(), str(model_path))
def test_load_model(self): devices = ["cpu"] if torch.cuda.is_available(): devices.append("cuda") for device in devices: device = torch.device(device) with self.subTest(f"Device {device}"): model = load_model(device) self.assertTrue(isinstance(model, torch.nn.Module))
def test_estimate(self): image = Image.open(Path(__file__).resolve().parent / 'datas' / 'sample.jpg') devices = ["cpu"] if torch.cuda.is_available(): devices.append("cuda") for device in devices: device = torch.device(device) model = load_model(device) with self.subTest(f"Device {device}"): label = estimate(data=image, model=model, device=device) self.assertTrue(0 <= label <= 9)
def main(): parser = argparse.ArgumentParser() parser.add_argument('--image_path', action='store', required=True, help='path to mnist image') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") # load model model = load_model(device, trained=True) # load image image = Image.open(args.image_path) label = estimate(image, model, device) print(f"Estimated label: {int(label)}")
estimator = EstimateAPI() return estimator.estimate() @app.route('/hc', methods=['GET']) def hc(): return 'alive' # global variables for ML from app.ml.model import load_model env_USE_GPU = os.getenv('USE_GPU') USE_GPU = not (env_USE_GPU is None or env_USE_GPU == '0') DEVICE = 'cuda' if USE_GPU else 'cpu' MODEL = load_model(device=DEVICE, trained=True) if __name__ == '__main__': config_file = Path( __file__).resolve().parents[1] / 'config' / 'api_config_local.yml' with config_file.open('r') as fp: config = yaml.load(fp) from logging import getLogger, StreamHandler, DEBUG logger = getLogger() handler = StreamHandler() handler.setLevel(DEBUG) logger.setLevel(DEBUG) logger.addHandler(handler)