def main(): train_loader = torch.utils.data.DataLoader(datasets.MNIST( './data', train=True, download=True, transform=transforms.Compose([transforms.ToTensor()])), batch_size=256, shuffle=True) test_loader = torch.utils.data.DataLoader(datasets.MNIST( './data', train=False, download=True, transform=transforms.Compose([ transforms.ToTensor(), ])), batch_size=9, shuffle=True) test_batch = None for x in test_loader: test_batch = x[0] break # Use cpu if no cuda available device = torch.device("cuda") #device = torch.device("cpu") model = LeNet().to(device) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9) train_with_logging(model, device, train_loader, optimizer, 200, 5, test_batch)
# 构建DataLoder train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True) valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE) # ============================ step 2.数据增强方法/5 模型 ============================ net = LeNet(classes=2) net.initialize_weights() # ============================ step 3/5 损失函数 ============================ criterion = nn.CrossEntropyLoss() # 选择损失函数 # ============================ step 4/5 优化器 ============================ optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9) # 选择优化器 scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # 设置学习率下降策略 # ============================ step 5/5 训练 ============================ train_curve = list() valid_curve = list() for epoch in range(MAX_EPOCH): loss_mean = 0. correct = 0. total = 0. net.train() for i, data in enumerate(train_loader):
def train_lenet(split_dir): """训练lenet""" set_seed() # 设置随机种子 rmb_label = {"1": 0, "100": 1} # 参数设置 MAX_EPOCH = 10 BATCH_SIZE = 16 LR = 0.01 log_interval = 10 val_interval = 1 """Step 1: 数据读取""" train_dir = os.path.join(split_dir, "train") valid_dir = os.path.join(split_dir, "valid") test_dir = os.path.join(split_dir, "test") norm_mean = [0.485, 0.456, 0.406] norm_std = [0.229, 0.224, 0.225] # 对训练数据进行变换,添加RandomCrop进行数据增强 train_transform = transforms.Compose([ transforms.Resize((32, 32)), transforms.RandomCrop(32, padding=4), transforms.ToTensor(), transforms.Normalize(norm_mean, norm_std), ]) # 对验证数据进行变换 valid_transform = transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(norm_mean, norm_std), ]) # 构建RMBDataset实例 train_data = RMBDataset(data_dir=train_dir, transform=train_transform) valid_data = RMBDataset(data_dir=valid_dir, transform=valid_transform) # 构建DataLoader train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True) valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE) """Step 2: 模型""" net = LeNet(classes=2) net.initialize_weights() """Step 3: 损失函数""" criterion = nn.CrossEntropyLoss() """Step 4: 优化器""" optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) """Step 5: 训练""" train_curve = [] valid_curve = [] figure_dir = os.path.join(os.path.abspath(os.path.dirname(__file__)), 'result') for epoch in range(MAX_EPOCH): loss_mean = 0. correct = 0. total = 0. net.train() for i, data in enumerate(train_loader): # forward inputs, labels = data outputs = net(inputs) # backward optimizer.zero_grad() loss = criterion(outputs, labels) loss.backward() # update weights optimizer.step() # 统计分类情况 _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).squeeze().sum().numpy() # 打印训练信息 loss_mean += loss.item() train_curve.append(loss.item()) if (i + 1) % log_interval == 0: loss_mean = loss_mean / log_interval print( "Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}" .format(epoch, MAX_EPOCH, i + 1, len(train_loader), loss_mean, correct / total)) loss_mean = 0. scheduler.step() # 更新学习率 # validate the model if (epoch + 1) % val_interval == 0: correct_val = 0. total_val = 0. loss_val = 0. net.eval() with torch.no_grad(): for j, data in enumerate(valid_loader): inputs, labels = data outputs = net(inputs) loss = criterion(outputs, labels) _, predicted = torch.max(outputs.data, 1) total_val += labels.size(0) correct_val += ( predicted == labels).squeeze().sum().numpy() loss_val += loss.item() valid_curve.append(loss_val / valid_loader.__len__()) print( "Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}" .format(epoch, MAX_EPOCH, j + 1, len(valid_loader), loss_val, correct_val / total_val)) train_x = range(len(train_curve)) train_y = train_curve train_iters = len(train_loader) valid_x = np.arange( 1, len(valid_curve) + 1 ) * train_iters * val_interval # 由于valid中记录的是epochloss,需要对记录点进行转换到iterations valid_y = valid_curve plt.plot(train_x, train_y, label='Train') plt.plot(valid_x, valid_y, label='Valid') plt.legend(loc='upper right') plt.ylabel('loss value') plt.xlabel('Iteration') figure_path = os.path.join(figure_dir, '0201.png') plt.savefig(figure_path) plt.close()
trainloader = th.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2) testset = tv.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = th.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) label_loss = nn.CrossEntropyLoss() t_opt = optim.SGD(teacher.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) # Train for epoch in range(200): epoch_teacher = 0.0 if epoch == 100 or epoch == 150: for group in t_opt.param_groups: group['lr'] *= 0.1 for X, y in trainloader: X, y = V(X.cuda()), V(y.cuda()) t_preds = teacher(X) t_loss = label_loss(t_preds, y) t_opt.zero_grad()
def main(): parser = argparse.ArgumentParser() mode_group = parser.add_mutually_exclusive_group(required=True) mode_group.add_argument("--train", action="store_true", help="To train the network.") mode_group.add_argument("--test", action="store_true", help="To test the network.") parser.add_argument("--epochs", default=10, type=int, help="Desired number of epochs.") parser.add_argument("--dropout", action="store_true", help="Whether to use dropout or not.") parser.add_argument("--uncertainty", action="store_true", help="Use uncertainty or not.") parser.add_argument("--dataset", action="store_true", help="The dataset to use.") parser.add_argument("--outsample", action="store_true", help="Use out of sample test image") uncertainty_type_group = parser.add_mutually_exclusive_group() uncertainty_type_group.add_argument( "--mse", action="store_true", help= "Set this argument when using uncertainty. Sets loss function to Expected Mean Square Error." ) uncertainty_type_group.add_argument( "--digamma", action="store_true", help= "Set this argument when using uncertainty. Sets loss function to Expected Cross Entropy." ) uncertainty_type_group.add_argument( "--log", action="store_true", help= "Set this argument when using uncertainty. Sets loss function to Negative Log of the Expected Likelihood." ) dataset_type_group = parser.add_mutually_exclusive_group() dataset_type_group.add_argument( "--mnist", action="store_true", help="Set this argument when using MNIST dataset") dataset_type_group.add_argument( "--emnist", action="store_true", help="Set this argument when using EMNIST dataset") dataset_type_group.add_argument( "--CIFAR", action="store_true", help="Set this argument when using CIFAR dataset") dataset_type_group.add_argument( "--fmnist", action="store_true", help="Set this argument when using FMNIST dataset") args = parser.parse_args() if args.dataset: if args.mnist: from mnist import dataloaders, label_list elif args.CIFAR: from CIFAR import dataloaders, label_list elif args.fmnist: from fashionMNIST import dataloaders, label_list if args.train: num_epochs = args.epochs use_uncertainty = args.uncertainty num_classes = 10 model = LeNet(dropout=args.dropout) if use_uncertainty: if args.digamma: criterion = edl_digamma_loss elif args.log: criterion = edl_log_loss elif args.mse: criterion = edl_mse_loss else: parser.error( "--uncertainty requires --mse, --log or --digamma.") else: criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.005) exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) device = get_device() model = model.to(device) model, metrics = train_model(model, dataloaders, num_classes, criterion, optimizer, scheduler=exp_lr_scheduler, num_epochs=num_epochs, device=device, uncertainty=use_uncertainty) state = { "epoch": num_epochs, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), } if use_uncertainty: if args.digamma: torch.save(state, "./results/model_uncertainty_digamma.pt") print("Saved: ./results/model_uncertainty_digamma.pt") if args.log: torch.save(state, "./results/model_uncertainty_log.pt") print("Saved: ./results/model_uncertainty_log.pt") if args.mse: torch.save(state, "./results/model_uncertainty_mse.pt") print("Saved: ./results/model_uncertainty_mse.pt") else: torch.save(state, "./results/model.pt") print("Saved: ./results/model.pt") elif args.test: use_uncertainty = args.uncertainty device = get_device() model = LeNet() model = model.to(device) optimizer = optim.Adam(model.parameters()) if use_uncertainty: if args.digamma: checkpoint = torch.load( "./results/model_uncertainty_digamma.pt") if args.log: checkpoint = torch.load("./results/model_uncertainty_log.pt") if args.mse: checkpoint = torch.load("./results/model_uncertainty_mse.pt") else: checkpoint = torch.load("./results/model.pt") filename = "./results/rotate.jpg" model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) model.eval() if args.outsample: img = Image.open("./data/arka.jpg").convert('L').resize((28, 28)) img = TF.to_tensor(img) img.unsqueeze_(0) else: a = iter(dataloaders['test']) img, label = next(a) rotating_image_classification(model, img, filename, label_list, uncertainty=use_uncertainty) img = transforms.ToPILImage()(img[0][0]) test_single_image(model, img, label_list, uncertainty=use_uncertainty)
net.conv2.register_forward_hook(get_activation('conv2')) # Log architecture dummy_input = torch.zeros(1, 1, net_input_size, net_input_size) writer.add_graph( net, input_to_model=dummy_input, # verbose=True, ) # Loss criterion = nn.CrossEntropyLoss() # Optimizer optimizer = optim.SGD( net.parameters(), lr=0.001, momentum=0.9, ) print('Training...') epochs = 41 log_every_n_epochs = 10 progress = tqdm(range(epochs)) iteration = 0 features_shape = len(dataset), net.fc2.out_features for epoch in progress: # loop over the dataset multiple times features = np.empty(features_shape) epoch_labels = [] # Log epoch validation loss
num_workers=args.worker ) val_loader = torch.utils.data.DataLoader( torchvision.datasets.ImageFolder(args.val, transform = data_transforms), batch_size=args.batch_size, shuffle=True, num_workers=args.worker ) return train_loader, val_loader train_loader, val_loader = get_dataloader() classes = 10 net = LeNet(classes) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # model to gpu net.to(device) print("Start Training...") os.makedirs('./expr', exist_ok=True) for epoch in range(1000): net.train() running_loss = 0.0 for i, data in enumerate(train_loader): inputs, labels = data inputs, labels = inputs.to(device, dtype=torch.float), labels.to(device) optimizer.zero_grad()
num_workers=args.num_workers) testset = torchvision.datasets.MNIST(root=args.data_path, train=False, download=True, transform=transform['test']) testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) net = LeNet() print(net) criterion = torch.nn.CrossEntropyLoss() optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=5e-4) exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=int(args.epochs / 2.5), gamma=0.1) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(device) net.to(device) max_precision = 0.0 for epoch in range(args.epochs): # loop over the dataset multiple times running_loss = 0.0 avg_loss = 0.0 k = 0 for i, data in enumerate(trainloader, 0): # get the inputs
import torch.nn as nn import torch as t t.set_num_threads(8) from lenet import LeNet net = LeNet() # loss and optim from torch import optim criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # training for epoch in range(2): running_loss = 0.0 for i, data in enumerate(trainloader, 0): inputs, labels = data # grad to zero optimizer.zero_grad() # forward + backward outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() # weight
def main(): # Training settings parser = argparse.ArgumentParser(description='PyTorch MNIST Example') 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=14, metavar='N', help='number of epochs to train (default: 14)') parser.add_argument('--lr', type=float, default=1.0, metavar='LR', help='learning rate (default: 1.0)') parser.add_argument('--gamma', type=float, default=0.7, metavar='M', help='Learning rate step gamma (default: 0.7)') parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training') parser.add_argument('--seed', type=int, default=1, metavar='S', help='random seed (default: 1)') parser.add_argument( '--log-interval', type=int, default=10, metavar='N', help='how many batches to wait before logging training status') parser.add_argument('--save-model', action='store_true', default=False, help='For Saving the current Model') args = parser.parse_args() use_cuda = not args.no_cuda and torch.cuda.is_available() torch.manual_seed(args.seed) device = torch.device("cuda" if use_cuda else "cpu") kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {} train_loader = torch.utils.data.DataLoader(datasets.MNIST( '../data', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ])), batch_size=args.batch_size, shuffle=True, **kwargs) test_loader = torch.utils.data.DataLoader(datasets.MNIST( '../data', train=False, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ])), batch_size=args.test_batch_size, shuffle=True, **kwargs) model = LeNet().to(device) optimizer = optim.Adadelta(model.parameters(), lr=args.lr) scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma) for epoch in range(1, args.epochs + 1): train(args, model, device, train_loader, optimizer, epoch) test(args, model, device, test_loader) scheduler.step() if args.save_model: model.export_mat()
class Ui(QtWidgets.QMainWindow): buttons = [ "load_image", "color_conversion", "image_flipping", "blending", "global_threshold", "local_threshold", "gaussian", "sobel_x", "sobel_y", "magnitude", "rst", "show_train_image", "show_hyper", "train_1", "pt", "inference", "ok", "show_train_result", "cancel"] inputs = ["angle", "scale", "tx", "ty", "test_index"] def __init__(self): super(Ui, self).__init__() uic.loadUi('main_window.ui', self) self.get_widgets() self.get_input() self.bind_event() self.param_setup() self.torch_setup() self.show() def get_widgets(self): for btn in self.buttons: setattr(self, btn, self.findChild(QtWidgets.QPushButton, btn)) def get_input(self): for inp in self.inputs: setattr(self, inp, self.findChild(QtWidgets.QLineEdit, inp)) def bind_event(self): for btn in self.buttons: getattr(self, btn).clicked.connect(partial( getattr(events, btn), self)) def param_setup(self): self.batch_size = 32 self.learning_rate = 0.001 self.opt = "SGD" self.loss_list = [] self.loss_epoch = [] self.acc_train_epoch = [] self.acc_test_epoch = [] self.compose = transforms.Compose([ transforms.Resize((32,32)), transforms.ToTensor() ]) def torch_setup(self): self.data_train = MNIST('./data/mnist', train=True, download=True, transform=self.compose) self.data_test = MNIST('./data/mnist', train=False, download=True, transform=self.compose) self.data_train_loader = DataLoader(self.data_train, batch_size=self.batch_size, shuffle=True, num_workers=4) self.data_test_loader = DataLoader(self.data_test, batch_size=self.batch_size, num_workers=4) self.criterion = nn.CrossEntropyLoss() self.net = LeNet() self.optimizer = getattr(optim, self.opt)(self.net.parameters(), lr=self.learning_rate) try: self.net.load_state_dict(load('model_params.pkl')) self.loaded = True print("Loaded") except Exception as e: print(e) self.loaded = False print("Not loaded") def train(self, epoch): self.net.train() self.loss_list = [] correct, total = 0, 0 for i, (images, labels) in enumerate(self.data_train_loader): self.optimizer.zero_grad() output = self.net(images) loss = self.criterion(output, labels) pred = output.data.max(1, keepdim=True)[1] correct += np.sum(np.squeeze(pred.eq(labels.data.view_as(pred))).cpu().numpy()) total += images.size(0) self.loss_list.append(loss.detach().cpu().item()) if i % 100 == 0: print(f'Train - Epoch {epoch}, Batch: {i}, Loss: {loss.detach().cpu().item()}') loss.backward() self.optimizer.step() self.acc_train_epoch.append(correct/total) self.loss_epoch.append(sum(self.loss_list)/len(self.loss_list)) def test(self): self.net.eval() total_correct, avg_loss = 0, 0.0 for i, (images, labels) in enumerate(self.data_test_loader): output = self.net(images) avg_loss += self.criterion(output, labels).sum() pred = output.detach().max(1)[1] total_correct += pred.eq(labels.view_as(pred)).sum() avg_loss /= len(self.data_test) acc = float(total_correct)/len(self.data_test) self.acc_test_epoch.append(acc) def test_and_train(self, epoch): self.train(epoch) self.test()
def training(model_name, trainloader, validloader, input_channel=3, epochs=1, resume=True, self_define=True, only_print=False): # load self defined or official net assert model_name in ["LeNet", "VGG16", "ResNet", "DenseNet"] if self_define: if model_name == "LeNet": net = LeNet(input_channel) elif model_name == "VGG16": net = VGG16(input_channel) elif model_name == "ResNet": net = ResNet(input_channel) elif model_name == "DenseNet": net = DenseNet(input_channel) else: if model_name == "LeNet": net = LeNet(input_channel) # on official LeNet elif model_name == "VGG16": net = models.vgg16_bn(pretrained=False, num_classes=10) elif model_name == "ResNet": net = models.resnet50(pretrained=False, num_classes=10) elif model_name == "DenseNet": net = models.DenseNet(num_classes=10) # sum of net parameters number print("Number of trainable parameters in %s : %f" % (model_name, sum(p.numel() for p in net.parameters() if p.requires_grad))) # print model structure if only_print: print(net) return # resume training param_path = "./model/%s_%s_parameter.pt" % (model_name, "define" if self_define else "official") if resume: if os.path.exists(param_path): net.load_state_dict(torch.load(param_path)) net.train() print("Resume training " + model_name) else: print("Train %s from scratch" % model_name) else: print("Train %s from scratch" % model_name) # define loss function and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # train on GPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print('train on %s' % device) net.to(device) running_loss = 0.0 train_losses = [] valid_losses = [] mini_batches = 125 * 5 for epoch in range(epochs): for i, data in enumerate(trainloader, 0): # get one batch # inputs, labels = data inputs, labels = data[0].to(device), data[1].to(device) # switch model to training mode, clear gradient accumulators net.train() optimizer.zero_grad() # forward + backward + optimize outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() # print statistics running_loss += loss.item() if i % mini_batches == mini_batches - 1: # print and valid every <mini_batches> mini-batches # validate model in validation dataset valid_loss = valid(net, validloader, criterion, device) print('[%d, %5d] train loss: %.3f, validset loss: %.3f' % ( epoch + 1, i + 1, running_loss / mini_batches, valid_loss)) train_losses.append(running_loss / mini_batches) valid_losses.append(valid_loss) running_loss = 0.0 # save parameters torch.save(net.state_dict(), param_path) # # save checkpoint # torch.save({ # 'epoch': epoch, # 'model_state_dict': net.state_dict(), # 'optimizer_state_dict': optimizer.state_dict(), # 'loss': loss # }, "./checkpoints/epoch_" + str(epoch) + ".tar") print('Finished Training, %d images in all' % (len(train_losses) * batch_size * mini_batches / epochs)) # draw loss curve assert len(train_losses) == len(valid_losses) loss_x = range(0, len(train_losses)) plt.plot(loss_x, train_losses, label="train loss") plt.plot(loss_x, valid_losses, label="valid loss") plt.title("Loss for every %d mini-batch" % mini_batches) plt.xlabel("%d mini-batches" % mini_batches) plt.ylabel("Loss") plt.legend() plt.savefig(model_name + "_loss.png") plt.show()
def main(): parser = argparse.ArgumentParser() mode_group = parser.add_mutually_exclusive_group(required=True) mode_group.add_argument("--train", action="store_true", help="To train the network.") mode_group.add_argument("--test", action="store_true", help="To test the network.") mode_group.add_argument("--examples", action="store_true", help="To example MNIST data.") parser.add_argument("--epochs", default=10, type=int, help="Desired number of epochs.") parser.add_argument("--dropout", action="store_true", help="Whether to use dropout or not.") parser.add_argument("--uncertainty", action="store_true", help="Use uncertainty or not.") uncertainty_type_group = parser.add_mutually_exclusive_group() uncertainty_type_group.add_argument( "--mse", action="store_true", help= "Set this argument when using uncertainty. Sets loss function to Expected Mean Square Error." ) uncertainty_type_group.add_argument( "--digamma", action="store_true", help= "Set this argument when using uncertainty. Sets loss function to Expected Cross Entropy." ) uncertainty_type_group.add_argument( "--log", action="store_true", help= "Set this argument when using uncertainty. Sets loss function to Negative Log of the Expected Likelihood." ) args = parser.parse_args() if args.examples: examples = enumerate(dataloaders["val"]) batch_idx, (example_data, example_targets) = next(examples) fig = plt.figure() for i in range(6): plt.subplot(2, 3, i + 1) plt.tight_layout() plt.imshow(example_data[i][0], cmap="gray", interpolation="none") plt.title("Ground Truth: {}".format(example_targets[i])) plt.xticks([]) plt.yticks([]) plt.savefig("./images/examples.jpg") elif args.train: num_epochs = args.epochs use_uncertainty = args.uncertainty num_classes = 10 model = LeNet(dropout=args.dropout) if use_uncertainty: if args.digamma: criterion = edl_digamma_loss elif args.log: criterion = edl_log_loss elif args.mse: criterion = edl_mse_loss else: parser.error( "--uncertainty requires --mse, --log or --digamma.") else: criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=0.005) exp_lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) device = get_device() model = model.to(device) model, metrics = train_model(model, dataloaders, num_classes, criterion, optimizer, scheduler=exp_lr_scheduler, num_epochs=num_epochs, device=device, uncertainty=use_uncertainty) state = { "epoch": num_epochs, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), } if use_uncertainty: if args.digamma: torch.save(state, "./results/model_uncertainty_digamma.pt") print("Saved: ./results/model_uncertainty_digamma.pt") if args.log: torch.save(state, "./results/model_uncertainty_log.pt") print("Saved: ./results/model_uncertainty_log.pt") if args.mse: torch.save(state, "./results/model_uncertainty_mse.pt") print("Saved: ./results/model_uncertainty_mse.pt") else: torch.save(state, "./results/model.pt") print("Saved: ./results/model.pt") elif args.test: use_uncertainty = args.uncertainty device = get_device() model = LeNet() model = model.to(device) optimizer = optim.Adam(model.parameters()) if use_uncertainty: if args.digamma: checkpoint = torch.load( "./results/model_uncertainty_digamma.pt") filename = "./results/rotate_uncertainty_digamma.jpg" if args.log: checkpoint = torch.load("./results/model_uncertainty_log.pt") filename = "./results/rotate_uncertainty_log.jpg" if args.mse: checkpoint = torch.load("./results/model_uncertainty_mse.pt") filename = "./results/rotate_uncertainty_mse.jpg" else: checkpoint = torch.load("./results/model.pt") filename = "./results/rotate.jpg" model.load_state_dict(checkpoint["model_state_dict"]) optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) model.eval() rotating_image_classification(model, digit_one, filename, uncertainty=use_uncertainty) img = Image.open("./data/one.jpg").convert('L') test_single_image(model, img, uncertainty=use_uncertainty)
def main(): global args, rank, world_size, best_prec1, dataset_len if args.dist == 1: rank, world_size = dist_init() else: rank = 0 world_size = 1 model = LeNet() model.cuda() param_copy = [ param.clone().type(torch.cuda.FloatTensor).detach() for param in model.parameters() ] for param in param_copy: param.requires_grad = True if args.dist == 1: model = DistModule(model) # define loss function (criterion) and optimizer criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(param_copy, args.base_lr, momentum=args.momentum, weight_decay=args.weight_decay) # optionally resume from a checkpoint last_iter = -1 # Data loading code train_dataset = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=False) val_dataset = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor(), download=False) dataset_len = len(train_dataset) args.max_iter = math.ceil( (dataset_len * args.epoch) / (world_size * args.batch_size)) if args.dist == 1: train_sampler = DistributedGivenIterationSampler(train_dataset, args.max_iter, args.batch_size, last_iter=last_iter) val_sampler = DistributedSampler(val_dataset, round_up=False) else: train_sampler = DistributedGivenIterationSampler(train_dataset, args.max_iter, args.batch_size, world_size=1, rank=0, last_iter=last_iter) val_sampler = None # pin_memory if true, will copy the tensor to cuda pinned memory train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=train_sampler) val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler) train(train_loader, val_loader, model, criterion, optimizer, param_copy)
numb_param({ "learning_rate": learning_rate, "epoch": epoch, "batch_size": batch_size }) # Register queue numb_queue(globals()) # it basically overrides variables in global scope def make_batch(in_data, targets, batch_size): for start in range(0, in_data.size()[0], batch_size): real_size = min((batch_size, in_data.size()[0] - start)) yield in_data[start:start + real_size], targets[start:start + real_size] criterion = nn.MSELoss() optimizer = optim.SGD(model.parameters(), lr=learning_rate) for epk in range(epoch): batches = make_batch(dummy_inputs, dummy_targets, batch_size) for train_exp, target in batches: model.zero_grad() output = model.forward(train_exp) loss = criterion.forward(output, target) loss.backward() optimizer.step() # print(model.state_dict()) print("Done Training!")
shuffle=True, num_workers=2) testset = tv.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test) testloader = th.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2) label_loss = nn.CrossEntropyLoss() distillation_loss = nn.MSELoss() sobolev = SobolevLoss(weight=0.1) s_opt = optim.SGD(student.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) t_opt = optim.SGD(teacher.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4) # Train for epoch in range(100): epoch_distillation = 0.0 epoch_teacher = 0.0 epoch_student = 0.0 epoch_sobolev = 0.0 for X, y in trainloader:
MNIST_train = torchvision.datasets.MNIST('./', download=True, train=True) MNIST_test = torchvision.datasets.MNIST('./', download=True, train=False) X_train, y_train = MNIST_train.data, MNIST_train.targets X_test, y_test = MNIST_test.data, MNIST_test.targets X_train, X_test = X_train.float(), X_test.float() X_train, X_test = X_train.unsqueeze(1), X_test.unsqueeze(1) model = LeNet() device = torch.device('cuda:0') if torch.cuda.is_available() else 'cpu' model = model.to(device) loss = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=1.0e-3) batch_size = 256 test_accuracy_history = [] test_loss_history = [] X_test = X_test.to(device) y_test = y_test.to(device) for epoch in range(1000): order = np.random.permutation(len(X_train)) for start_index in range(0, len(X_train), batch_size): optimizer.zero_grad()