def deploy(args, data_loader): model = Network(k=args.network_k, att_type=args.network_att_type, kernel3=args.kernel3, width=args.network_width, dropout=args.network_dropout, compensate=True, norm=args.norm, inp_channels=args.input_channels) print(model) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device) checkpoint_path = os.path.join(args.logdir, 'best_checkpoint.pth') if os.path.isfile(checkpoint_path): checkpoint = torch.load(checkpoint_path) model.load_state_dict(checkpoint['state_dict']) else: raise Exception('Couldnt load checkpoint.') df = pd.DataFrame(columns=['img', 'label', 'pred']) with tqdm(enumerate(data_loader)) as pbar: for i, (images, labels) in pbar: raw_label = labels raw_images = images if torch.cuda.is_available(): images = images.cuda() labels = labels.cuda() images.requires_grad = True # Forward pass outputs, att, localised = model(images, True) localised = F.softmax(localised.data, 3)[..., 1] predicted = torch.argmax(outputs.data, 1) saliency = torch.autograd.grad(outputs[:, 1].sum(), images)[0].data localised = localised[0].cpu().numpy() saliency = torch.sqrt((saliency[0]**2).mean(0)).cpu().numpy() raw_img = np.transpose(raw_images.numpy(), (0, 2, 3, 1)).squeeze() np.save(os.path.join(args.outpath, 'pred_{}.npy'.format(i)), localised) np.save(os.path.join(args.outpath, 'sal_{}.npy'.format(i)), saliency) df.loc[len(df)] = [ i, raw_label.numpy().squeeze(), predicted.cpu().numpy().squeeze() ] df.to_csv(os.path.join(args.outpath, 'pred.csv'), index=False) print('done - stopping now')
def main(params): print("Loading dataset ... ") with open(params['train_data_pkl'], 'rb') as f: train_data = pkl.load(f) with open(params['train_anno_pkl'], 'rb') as f: train_anno = pkl.load(f) """ with open(params['val_data_pkl'], 'rb') as f: val_data = pkl.load(f) with open(params['val_anno_pkl'], 'rb') as f: val_anno = pkl.load(f) """ # Train dataset and Train dataloader train_data = np.transpose(train_data, (0, 3, 1, 2)) train_dataset = torch.utils.data.TensorDataset( torch.FloatTensor(train_data), torch.LongTensor(train_anno)) train_loader = dataloader.DataLoader(train_dataset, params['batch_size'], shuffle=True, collate_fn=collate_fn) """ # Validation dataset and Validation dataloader val_data = np.transpose(val_data, (0, 3, 1, 2)) val_dataset = torch.utils.data.TensorDataset( torch.FloatTensor(val_data), torch.LongTensor(val_anno)) val_loader = dataloader.DataLoader( val_dataset, params['batch_size'], collate_fn=collate_fn) """ # the number of layers in each dense block n_layers_list = [4, 5, 7, 10, 12, 15, 12, 10, 7, 5, 4] print("Constructing the network ... ") # Define the network densenet = Network(n_layers_list, 5).to(device) if os.path.isfile(params['model_from']): print("Starting from the saved model") densenet.load_state_dict(torch.load(params['model_from'])) else: print("Couldn't find the saved model") print("Starting from the bottom") print("Training the model ...") # hyperparameter, optimizer, criterion learning_rate = params['lr'] optimizer = torch.optim.RMSprop(densenet.parameters(), learning_rate, weight_decay=params['l2_reg']) criterion = nn.CrossEntropyLoss() for epoch in range(params['max_epoch']): for i, (img, label) in enumerate(train_loader): img = img.to(device) label = label.to(device) # forward-propagation pred = densenet(img) # flatten for all pixel pred = pred.view((-1, params['num_answers'])) label = label.view((-1)) # get loss loss = criterion(pred, label) # back-propagation optimizer.zero_grad() loss.backward() optimizer.step() print("Epoch: %d, Steps:[%d/%d], Loss: %.4f" % (epoch, i, len(train_loader), loss.data)) learning_rate *= 0.995 optimizer = torch.optim.RMSprop(densenet.parameters(), learning_rate, weight_decay=params['l2_reg']) if (epoch + 1) % 10 == 0: print("Saved the model") torch.save(densenet.state_dict(), params['model_save'])