def pred_prob(arg_path, field_path, pth_path, doc, device=torch.device('cpu')): # Load args # with open(arg_path) as f: # args = json.load(f)['args'] arg_path = os.path.join( 'https://raw.githubusercontent.com/qianyingw/rob-pome/master/rob-app', arg_path) with urllib.request.urlopen(arg_path) as url: args = json.loads(url.read().decode())['args'] # Load TEXT field field_url = os.path.join( 'https://github.com/qianyingw/rob-pome/raw/master/rob-app', field_path) field_path = wget.download(field_url) with open(field_path, "rb") as fin: TEXT = dill.load(fin) os.remove(field_path) unk_idx = TEXT.vocab.stoi[TEXT.unk_token] # 0 pad_idx = TEXT.vocab.stoi[TEXT.pad_token] # 1 # Load model if args['net_type'] == 'cnn': sizes = args['filter_sizes'].split(',') sizes = [int(s) for s in sizes] model = ConvNet(vocab_size=args['max_vocab_size'] + 2, embedding_dim=args['embed_dim'], n_filters=args['num_filters'], filter_sizes=sizes, output_dim=2, dropout=args['dropout'], pad_idx=pad_idx, embed_trainable=args['embed_trainable'], batch_norm=args['batch_norm']) if args['net_type'] == 'attn': model = AttnNet(vocab_size=args['max_vocab_size'] + 2, embedding_dim=args['embed_dim'], rnn_hidden_dim=args['rnn_hidden_dim'], rnn_num_layers=args['rnn_num_layers'], output_dim=2, bidirection=args['bidirection'], rnn_cell_type=args['rnn_cell_type'], dropout=args['dropout'], pad_idx=pad_idx, embed_trainable=args['embed_trainable'], batch_norm=args['batch_norm'], output_attn=False) # Load checkpoint pth_url = os.path.join( 'https://github.com/qianyingw/rob-pome/raw/master/rob-app', pth_path) pth_path = wget.download(pth_url) checkpoint = torch.load(pth_path, map_location=device) os.remove(pth_path) state_dict = checkpoint['state_dict'] model.load_state_dict(state_dict, strict=False) model.cpu() # Load pre-trained embedding pretrained_embeddings = TEXT.vocab.vectors model.embedding.weight.data.copy_(pretrained_embeddings) model.embedding.weight.data[unk_idx] = torch.zeros( args['embed_dim']) # Zero the initial weights for <unk> tokens model.embedding.weight.data[pad_idx] = torch.zeros( args['embed_dim']) # Zero the initial weights for <pad> tokens # Tokenization tokens = [tok.text.lower() for tok in nlp.tokenizer(doc)] idx = [TEXT.vocab.stoi[t] for t in tokens] while len(idx) < args['max_token_len']: idx = idx + [1] * args['max_token_len'] if len(idx) > args['max_token_len']: idx = idx[:args['max_token_len']] # Prediction model.eval() doc_tensor = torch.LongTensor(idx).to(device) doc_tensor = doc_tensor.unsqueeze( 1) # bec AttnNet input shape is [seq_len, batch_size] probs = model(doc_tensor) probs = probs.data.cpu().numpy()[0] # print("Prob of RoB reported: {:.4f}".format(probs[1])) return probs[1]
# zero the parameter gradients optimizer.zero_grad() input_patch = input_patch.to(device) gt_patch = gt_patch.to(device) # forward + backward + optimize output = net(input_patch) loss = criterion(output, gt_patch) loss.backward() optimizer.step() running_loss += loss.item() if cnt == (num_ids - 1): print("epoch: %d, loss: %.4f" % (epoch, running_loss / num_ids)) if epoch % save_freq == 0: if not os.path.isdir(result_dir + '%04d' % epoch): os.makedirs(result_dir + '%04d' % epoch) temp = np.concatenate((gt_patch.cpu().detach().numpy()[0, :, :, :], output.cpu().detach().numpy()[0, :, :, :]), axis=1) scipy.misc.toimage(temp * 255, high=255, low=0, cmin=0, cmax=255).save(result_dir + '%04d/%05d_00_train_%d.jpg' % (epoch, train_id, ratio)) torch.save(net.cpu(), checkpoint_dir + 'sid_torchversion_model.ckpt')