def handler(context): dataset_alias = context.datasets dataset_id = dataset_alias['data'] # set alias specified in console data = list(load_dataset_from_api(dataset_id)) # data = list(load_dataset_from_path()) random.shuffle(data) x = np.array([_[0] for _ in data]) y = np.array([_[1] for _ in data]) train_size = int(len(x) * 0.7) x_train, x_test = x[:train_size], x[train_size:] y_train, y_test = y[:train_size], y[train_size:] nb_channels = 3 if K.image_data_format() == 'channels_first': x_train = x_train.reshape(x_train.shape[0], nb_channels, img_rows, img_cols) x_test = x_test.reshape(x_test.shape[0], nb_channels, img_rows, img_cols) input_shape = (nb_channels, img_rows, img_cols) else: x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, nb_channels) x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, nb_channels) input_shape = (img_rows, img_cols, nb_channels) x_train = x_train.astype('float32') x_test = x_test.astype('float32') print('x_train shape:', x_train.shape) print(x_train.shape[0], 'train samples') print(x_test.shape[0], 'test samples') # convert class vectors to binary class matrices y_train = keras.utils.to_categorical(y_train, num_classes) y_test = keras.utils.to_categorical(y_test, num_classes) model = create_model(input_shape) tensorboard = TensorBoard(log_dir=log_path, histogram_freq=0, write_graph=True, write_images=False) statistics = Statistics() # Do you want to add `checkpoint` to callback as well? model.compile(loss=keras.losses.categorical_crossentropy, optimizer=keras.optimizers.Adam(lr=LEARNING_RATE), metrics=['accuracy']) # fit_generator # image loader model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, verbose=1, validation_data=(x_test, y_test), callbacks=[tensorboard, statistics]) score = model.evaluate(x_test, y_test, verbose=0) model.save(os.path.join(ABEJA_TRAINING_RESULT_DIR, 'model.h5')) print('Test loss:', score[0]) print('Test accuracy:', score[1])
def handler(context): dataset_alias = context.datasets dataset_id = dataset_alias['train'] # set alias specified in console data = list(load_dataset_from_api(dataset_id)) np.random.seed(0) data = np.random.permutation(data) nb_data = len(data) nb_train = int(7 * nb_data // 10) train_data_raw = data[:nb_train] test_data_raw = data[nb_train:] simple_net = SimpleNet(num_classes) model = L.Classifier(simple_net) if USE_GPU >= 0: chainer.cuda.get_device(USE_GPU).use() # Make a specified GPU current model.to_gpu() def make_optimizer(model, alpha=0.001, beta1=0.9): optimizer = chainer.optimizers.Adam(alpha=alpha, beta1=beta1) optimizer.setup(model) return optimizer optimizer = make_optimizer(model) train_data = ImageDatasetFromAPI(train_data_raw) train_iter = chainer.iterators.SerialIterator(train_data, batch_size) test_data = ImageDatasetFromAPI(test_data_raw) test_iter = chainer.iterators.SerialIterator(test_data, batch_size, repeat=False, shuffle=False) # Set up a trainer updater = training.StandardUpdater(train_iter, optimizer, device=USE_GPU) trainer = training.Trainer(updater, (epochs, 'epoch'), out=ABEJA_TRAINING_RESULT_DIR) trainer.extend(extensions.Evaluator(test_iter, model, device=USE_GPU)) trainer.extend(extensions.snapshot_object(simple_net, 'simple_net.model'), trigger=(epochs, 'epoch')) report_entries = [ 'epoch', 'main/loss', 'validation/main/loss', 'main/accuracy', 'validation/main/accuracy' ] trainer.extend(extensions.LogReport()) trainer.extend(Statistics(report_entries, epochs), trigger=(1, 'epoch')) trainer.extend(Tensorboard(report_entries, out_dir=log_path)) trainer.extend(extensions.PrintReport(report_entries)) trainer.run()
def handler(context): dataset_alias = context.datasets trainval_2007_dataset_id = dataset_alias['trainval2007'] trainval_2012_dataset_id = dataset_alias['trainval2012'] test_2007_dataset_id = dataset_alias['test2007'] trainval_2007_dataset = list( load_dataset_from_api(trainval_2007_dataset_id)) trainval_2012_dataset = list( load_dataset_from_api(trainval_2012_dataset_id)) test_2007_dataset = list(load_dataset_from_api(test_2007_dataset_id)) if network_model == 'ssd300': model = SSD300(n_fg_class=len(voc_bbox_label_names), pretrained_model='imagenet') elif network_model == 'ssd512': model = SSD512(n_fg_class=len(voc_bbox_label_names), pretrained_model='imagenet') model.use_preset('evaluate') train_chain = MultiboxTrainChain(model) if USE_GPU >= 0: chainer.cuda.get_device_from_id(USE_GPU).use() model.to_gpu() trainval_2007 = DetectionDatasetFromAPI(trainval_2007_dataset) trainval_2012 = DetectionDatasetFromAPI(trainval_2012_dataset) test_2007 = DetectionDatasetFromAPI(test_2007_dataset, use_difficult=True, return_difficult=True) train = TransformDataset(ConcatenatedDataset(trainval_2007, trainval_2012), Transform(model.coder, model.insize, model.mean)) train_iter = chainer.iterators.SerialIterator(train, BATCHSIZE) test_iter = chainer.iterators.SerialIterator(test_2007, BATCHSIZE, repeat=False, shuffle=False) # initial lr is set to 1e-3 by ExponentialShift optimizer = chainer.optimizers.MomentumSGD() optimizer.setup(train_chain) for param in train_chain.params(): if param.name == 'b': param.update_rule.add_hook(GradientScaling(2)) else: param.update_rule.add_hook(WeightDecay(0.0005)) updater = training.updaters.StandardUpdater(train_iter, optimizer, device=USE_GPU) trainer = training.Trainer(updater, (nb_iterations, 'iteration'), out=ABEJA_TRAINING_RESULT_DIR) trainer.extend(extensions.ExponentialShift('lr', 0.1, init=1e-3), trigger=triggers.ManualScheduleTrigger([80000, 100000], 'iteration')) trainer.extend(DetectionVOCEvaluator(test_iter, model, use_07_metric=True, label_names=voc_bbox_label_names), trigger=(10000, 'iteration')) log_interval = 100, 'iteration' trainer.extend(extensions.LogReport(trigger=log_interval)) trainer.extend(extensions.observe_lr(), trigger=log_interval) print_entries = [ 'iteration', 'main/loss', 'main/loss/loc', 'main/loss/conf', 'validation/main/map' ] report_entries = [ 'epoch', 'iteration', 'lr', 'main/loss', 'main/loss/loc', 'main/loss/conf', 'validation/main/map' ] trainer.extend(Statistics(report_entries, nb_iterations, obs_key='iteration'), trigger=log_interval) trainer.extend(Tensorboard(report_entries, out_dir=log_path)) trainer.extend(extensions.PrintReport(print_entries), trigger=log_interval) trainer.extend(extensions.snapshot_object( model, 'model_iter_{.updater.iteration}'), trigger=(nb_iterations, 'iteration')) trainer.run()
def handler(context): # Data from ABEJA Platform dataset_alias = context.datasets dataset_id = dataset_alias['train'] data = list(load_dataset_from_api(dataset_id)) # Divide data to train adn validation np.random.seed(0) data = np.random.permutation(data) nb_data = len(data) nb_train = int(7 * nb_data // 10) train_data_raw = data[:nb_train] test_data_raw = data[nb_train:] # GPU use_cuda = USE_GPU >= 0 and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") net = SimpleNet(num_classes).to(device) criterion = torch.nn.CrossEntropyLoss().to(device) # optimizer optimizer = torch.optim.Adam(net.parameters()) # dataset kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} train_data = ImageDatasetFromAPI(train_data_raw, transform=transforms.Compose( [Preprocess(train=True), ToTensor()])) trainloader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=True, **kwargs) test_data = ImageDatasetFromAPI(test_data_raw, transform=transforms.Compose( [Preprocess(train=False), ToTensor()])) testloader = torch.utils.data.DataLoader(test_data, batch_size=batch_size, shuffle=True, **kwargs) # Reporter for ABEJA Platform statistics = Statistics(epochs) for epoch in range(epochs): net.train() train_total_loss = 0 train_total_acc = 0 train_total = 0 for i, (inputs, labels) in enumerate(trainloader): inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() _, predicted = torch.max(outputs.data, 1) train_total += labels.size(0) train_total_acc += (predicted == labels).sum().item() train_total_loss += loss.item() train_total_loss /= train_total train_total_acc /= train_total net.eval() test_total_loss = 0 test_total_acc = 0 test_total = 0 with torch.no_grad(): for i, (inputs, labels) in enumerate(testloader): inputs, labels = inputs.to(device), labels.to(device) outputs = net(inputs) loss = criterion(outputs, labels) _, predicted = torch.max(outputs.data, 1) test_total += labels.size(0) test_total_acc += (predicted == labels).sum().item() test_total_loss += loss.item() test_total_loss /= test_total test_total_acc /= test_total print( '[{:d}] main/loss: {:.3f} main/acc: {:.3f}, main/validation/loss: {:.3f}, main/validation/acc: {:.3f}' .format(epoch + 1, train_total_loss, train_total_acc, test_total_loss, test_total_acc)) statistics(epoch + 1, train_total_loss, train_total_acc, test_total_loss, test_total_acc) writer.add_scalar('main/loss', train_total_loss, epoch + 1) writer.add_scalar('main/acc', train_total_acc, epoch + 1) writer.add_scalar('main/validation/loss', test_total_loss, epoch + 1) writer.add_scalar('main/validation/acc', test_total_acc, epoch + 1) torch.save(net.state_dict(), os.path.join(ABEJA_TRAINING_RESULT_DIR, 'model.pth'))
USE_GPU = int(os.environ.get('USE_GPU', '-1')) log_path = os.path.join(ABEJA_TRAINING_RESULT_DIR, 'logs') writer = SummaryWriter(log_dir=log_path) use_cuda = USE_GPU >= 0 and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") batch_size = 32 lr = 1e-3 num_classes = 21 lr_steps = (80000, 100000, 120000) max_iter = 120000 min_dim = 300 statistics = Statistics(max_iter) def download(url, filename): import urllib urllib.request.urlretrieve(url, filename) def adjust_learning_rate(optimizer, gamma, step): """Sets the learning rate to the initial LR decayed by 10 at every specified step # Adapted from PyTorch Imagenet example: # https://github.com/pytorch/examples/blob/master/imagenet/main.py """ lr = lr * (gamma**(step)) for param_group in optimizer.param_groups:
def handler(context): # Triggers log_trigger = (50, 'iteration') validation_trigger = (2000, 'iteration') end_trigger = (nb_iterations, 'iteration') # Dataset dataset_alias = context.datasets train_dataset_id = dataset_alias['train'] val_dataset_id = dataset_alias['val'] train = SegmentationDatasetFromAPI(train_dataset_id) val = SegmentationDatasetFromAPI(val_dataset_id) class_weight = calc_weight(train) print(class_weight) train = TransformDataset(train, transform) # Iterator train_iter = iterators.SerialIterator(train, BATCHSIZE) val_iter = iterators.SerialIterator(val, BATCHSIZE, shuffle=False, repeat=False) # Model model = SegNetBasic(n_class=len(camvid_label_names)) model = PixelwiseSoftmaxClassifier(model, class_weight=class_weight) if USE_GPU >= 0: # Make a specified GPU current chainer.cuda.get_device_from_id(USE_GPU).use() model.to_gpu() # Copy the model to the GPU # Optimizer optimizer = optimizers.MomentumSGD(lr=0.1, momentum=0.9) optimizer.setup(model) optimizer.add_hook(chainer.optimizer_hooks.WeightDecay(rate=0.0005)) # Updater updater = training.updaters.StandardUpdater(train_iter, optimizer, device=USE_GPU) # Trainer trainer = training.Trainer(updater, end_trigger, out=ABEJA_TRAINING_RESULT_DIR) trainer.extend(extensions.LogReport(trigger=log_trigger)) trainer.extend(extensions.observe_lr(), trigger=log_trigger) trainer.extend(extensions.dump_graph('main/loss')) trainer.extend(extensions.snapshot_object( model.predictor, filename='model_iteration-{.updater.iteration}'), trigger=end_trigger) print_entries = [ 'iteration', 'main/loss', 'validation/main/miou', 'validation/main/mean_class_accuracy', 'validation/main/pixel_accuracy' ] report_entries = [ 'epoch', 'iteration', 'lr', 'main/loss', 'validation/main/miou', 'validation/main/mean_class_accuracy', 'validation/main/pixel_accuracy' ] trainer.extend(Statistics(report_entries, nb_iterations, obs_key='iteration'), trigger=log_trigger) trainer.extend(Tensorboard(report_entries, out_dir=log_path)) trainer.extend(extensions.PrintReport(print_entries), trigger=log_trigger) trainer.extend(SemanticSegmentationEvaluator(val_iter, model.predictor, camvid_label_names), trigger=validation_trigger) trainer.run()
def handler(context): class_labels = 10 dataset_alias = context.datasets train_dataset_id = dataset_alias['train'] test_dataset_id = dataset_alias['test'] train_data = list(load_dataset_from_api(train_dataset_id)) test_data = list(load_dataset_from_api(test_dataset_id)) train = ImageDatasetFromAPI(train_data, train=True) test = ImageDatasetFromAPI(test_data) net = utils.VGG.VGG(class_labels) model = L.Classifier(net) if USE_GPU >= 0: # Make a specified GPU current chainer.backends.cuda.get_device_from_id(USE_GPU).use() model.to_gpu() # Copy the model to the GPU optimizer = chainer.optimizers.MomentumSGD(learnrate) optimizer.setup(model) optimizer.add_hook(chainer.optimizer_hooks.WeightDecay(5e-4)) train_iter = chainer.iterators.SerialIterator(train, batchsize) test_iter = chainer.iterators.SerialIterator(test, batchsize, repeat=False, shuffle=False) stop_trigger = (epochs, 'epoch') # Early stopping option if early_stopping: stop_trigger = triggers.EarlyStoppingTrigger(monitor=early_stopping, verbose=True, max_trigger=(epochs, 'epoch')) # Set up a trainer updater = training.updaters.StandardUpdater(train_iter, optimizer, device=USE_GPU) trainer = training.Trainer(updater, stop_trigger, out=ABEJA_TRAINING_RESULT_DIR) # Evaluate the model with the test dataset for each epoch trainer.extend(extensions.Evaluator(test_iter, model, device=USE_GPU)) # Reduce the learning rate by half every 25 epochs. trainer.extend(extensions.ExponentialShift('lr', 0.5), trigger=(25, 'epoch')) # Take a snapshot at each epoch trainer.extend(extensions.snapshot_object(net, 'net.model'), trigger=(epochs, 'epoch')) # Write a log of evaluation statistics for each epoch trainer.extend(extensions.LogReport()) # Print selected entries of the log to stdout # Here "main" refers to the target link of the "main" optimizer again, and # "validation" refers to the default name of the Evaluator extension. # Entries other than 'epoch' are reported by the Classifier link, called by # either the updater or the evaluator. report_entries = [ 'epoch', 'main/loss', 'validation/main/loss', 'main/accuracy', 'validation/main/accuracy' ] trainer.extend(Statistics(report_entries, epochs), trigger=(1, 'epoch')) trainer.extend(Tensorboard(report_entries, out_dir=log_path)) trainer.extend(extensions.PrintReport(report_entries)) trainer.run()
from ssd import SSD from ssd import MultiBoxLoss from tools import download from utils.callbacks import Statistics if Parameters.RANDOM_SEED is not None: torch.manual_seed(Parameters.RANDOM_SEED) np.random.seed(Parameters.RANDOM_SEED) random.seed(Parameters.RANDOM_SEED) use_cuda = torch.cuda.is_available() # NOTE: multi gpu not supported yet device = torch.device("cuda" if use_cuda else "cpu") print("device: ", device) statistics = Statistics(Parameters.EPOCHS) log_path = os.path.join(Parameters.ABEJA_TRAINING_RESULT_DIR, 'logs') writer = SummaryWriter(log_dir=log_path) def handler(context): print( f'start training with parameters : {Parameters.as_dict()}, context : {context}' ) try: dataset_alias = context.datasets # for older version except AttributeError: dataset_alias = context['datasets']
def handler(context): dataset_alias = context.datasets data = list(load_dataset_from_api(dataset_alias['train'])) np.random.seed(0) data = np.random.permutation(data) nb_data = len(data) nb_train = int(7 * nb_data // 10) train_data_raw = data[:nb_train] test_data_raw = data[nb_train:] premodel = SSD300(n_fg_class=20, pretrained_model='voc0712') model = SSD300(n_fg_class=1) copy_ssd(model, premodel) model.use_preset('evaluate') train_chain = MultiboxTrainChain(model) if USE_GPU >= 0: chainer.cuda.get_device_from_id(USE_GPU).use() model.to_gpu() # initial lr is set to 1e-3 by ExponentialShift optimizer = chainer.optimizers.MomentumSGD() optimizer.setup(train_chain) for param in train_chain.params(): if param.name == 'b': param.update_rule.add_hook(GradientScaling(2)) else: param.update_rule.add_hook(WeightDecay(0.0005)) fix_ssd(train_chain) train_data = DetectionDatasetFromAPI(train_data_raw) test_data = DetectionDatasetFromAPI(test_data_raw, use_difficult=True, return_difficult=True) train_data = TransformDataset( train_data, Transform(model.coder, model.insize, model.mean)) train_iter = chainer.iterators.SerialIterator(train_data, BATCHSIZE) test_iter = chainer.iterators.SerialIterator(test_data, BATCHSIZE, repeat=False, shuffle=False) updater = training.updaters.StandardUpdater(train_iter, optimizer, device=USE_GPU) trainer = training.Trainer(updater, (nb_epochs, 'epoch'), out=ABEJA_TRAINING_RESULT_DIR) trainer.extend(extensions.ExponentialShift('lr', 0.1, init=1e-3), trigger=triggers.ManualScheduleTrigger([1200, 1600], 'epoch')) trainer.extend(DetectionVOCEvaluator(test_iter, model, use_07_metric=True, label_names=['cup']), trigger=(1, 'epoch')) log_interval = 1, 'epoch' trainer.extend(extensions.LogReport(trigger=log_interval)) print_entries = [ 'epoch', 'main/loss', 'main/loss/loc', 'main/loss/conf', 'validation/main/map' ] report_entries = [ 'epoch', 'lr', 'main/loss', 'main/loss/loc', 'main/loss/conf', 'validation/main/map' ] trainer.extend(Statistics(report_entries, nb_epochs), trigger=log_interval) trainer.extend(Tensorboard(report_entries, out_dir=log_path)) trainer.extend(extensions.PrintReport(print_entries), trigger=log_interval) trainer.extend(extensions.snapshot_object(model, 'model_epoch_{.updater.epoch}'), trigger=(nb_epochs, 'epoch')) trainer.run()
def handler(context): dataset_alias = context.datasets train_dataset_id = dataset_alias['train'] test_dataset_id = dataset_alias['test'] train_data = list(load_dataset_from_api(train_dataset_id)) test_data = list(load_dataset_from_api(test_dataset_id)) # Data print('==> Preparing data..') transform_train = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) transform_test = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)), ]) #trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train) trainset = ImageDatasetFromAPI(train_data, transform=transform_train) trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2) #testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testset = ImageDatasetFromAPI(test_data, transform=transform_test) testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) # Model print('==> Building model..') if model == 'VGG19': net = VGG('VGG19') elif model == 'ResNet18': net = ResNet18() elif model == 'PreActResNet18': net = PreActResNet18() elif model == 'GoogLeNet': net = GoogLeNet() elif model == 'DenseNet121': net = DenseNet121() elif model == 'ResNeXt29_2x64d': net = ResNeXt29_2x64d() elif model == 'MobileNet': net = MobileNet() elif model == 'MobileNetV2': net = MobileNetV2() elif model == 'DPN92': net = DPN92() elif model == 'ShuffleNetG2': net = ShuffleNetG2() elif model == 'SENet18': net = SENet18() elif model == 'ShuffleNetV2': net = ShuffleNetV2(1) net = net.to(device) #if device == 'cuda': # net = torch.nn.DataParallel(net) # cudnn.benchmark = True criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[150, 250], gamma=0.1) statistics = Statistics(epochs) for epoch in range(epochs): scheduler.step() train_loss, train_acc = train(net, optimizer, trainloader, criterion, epoch) test_loss, test_acc = test(net, testloader, criterion, epoch) # Reporting print( '[{:d}] main/loss: {:.3f} main/acc: {:.3f}, main/validation/loss: {:.3f}, main/validation/acc: {:.3f}' .format(epoch + 1, train_loss, train_acc, test_loss, test_acc)) statistics(epoch + 1, train_loss, train_acc, test_loss, test_acc) writer.add_scalar('main/loss', train_loss, epoch + 1) writer.add_scalar('main/acc', train_acc, epoch + 1) writer.add_scalar('main/validation/loss', test_loss, epoch + 1) writer.add_scalar('main/validation/acc', test_acc, epoch + 1) torch.save(net.state_dict(), os.path.join(ABEJA_TRAINING_RESULT_DIR, 'model.pth'))
def handler(context): # Dataset dataset_alias = context.datasets train_dataset_id = dataset_alias['train'] val_dataset_id = dataset_alias['val'] trainset = SegmentationDatasetFromAPI(train_dataset_id, transform=SegNetAugmentation(MEANS)) valset = SegmentationDatasetFromAPI(val_dataset_id, transform=SegNetAugmentation( MEANS, False)) class_weight = calc_weight( SegmentationDatasetFromAPI(train_dataset_id, transform=SegNetAugmentation(MEANS, False))) class_weight = class_weight.to(device) trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCHSIZE, shuffle=True, num_workers=0) valloader = torch.utils.data.DataLoader(valset, batch_size=BATCHSIZE, shuffle=False, num_workers=0) # Model net = SegNet(3, n_class=len(camvid_label_names)) net = net.to(device) # Optimizer #criterion = PixelwiseSoftmaxClassifier(weight=class_weight) criterion = torch.nn.CrossEntropyLoss(weight=class_weight, ignore_index=-1) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[150, 250], gamma=0.1) statistics = Statistics(epochs) for epoch in range(epochs): scheduler.step() train_loss, train_acc = train(net, optimizer, trainloader, criterion, epoch) test_loss, test_acc = test(net, valloader, criterion, epoch) # Reporting print( '[{:d}] main/loss: {:.3f} main/acc: {:.3f}, main/validation/loss: {:.3f}, main/validation/acc: {:.3f}' .format(epoch + 1, train_loss, train_acc, test_loss, test_acc)) statistics(epoch + 1, train_loss, train_acc, test_loss, test_acc) writer.add_scalar('main/loss', train_loss, epoch + 1) writer.add_scalar('main/acc', train_acc, epoch + 1) writer.add_scalar('main/validation/loss', test_loss, epoch + 1) writer.add_scalar('main/validation/acc', test_acc, epoch + 1) torch.save(net.state_dict(), os.path.join(ABEJA_TRAINING_RESULT_DIR, 'model.pth'))