Update app.py

app.py

@@ -16,17 +16,17 @@ parser = argparse.ArgumentParser(description='Image Captioning inference script'
 parser.add_argument('--max_seq_len', default=60, type=int, help='max sequence length')
 
 # Model parameters
-parser.add_argument('--height', default=32, type=int, metavar='N', help='image height')
-parser.add_argument('--width', default=32, type=int, metavar='N', help='image width')
+parser.add_argument('--height', default=224, type=int, metavar='N', help='image height')
+parser.add_argument('--width', default=224, type=int, metavar='N', help='image width')
 parser.add_argument('--channel', default=3, type=int, help='disable cuda')
-parser.add_argument('--enc_heads', default=12, type=int, help='number of encoder  heads')
-parser.add_argument('--enc_depth', default=9, type=int, help='number of encoder blocks')
-parser.add_argument('--dec_heads', default=12, type=int, help='number of decoder  heads')
+parser.add_argument('--enc_heads', default=8, type=int, help='number of encoder  heads')
+parser.add_argument('--enc_depth', default=4, type=int, help='number of encoder blocks')
+parser.add_argument('--dec_heads', default=8, type=int, help='number of decoder  heads')
 parser.add_argument('--dec_depth', default=1, type=int, help='number of decoder blocks')
-parser.add_argument('--patch_size', default=4, type=int, help='patch size')
-parser.add_argument('--dim', default=192, type=int, help='embedding dim of patch')
-parser.add_argument('--enc_mlp_dim', default=384, type=int, help='feed forward hidden_dim for an encoder block')
-parser.add_argument('--dec_mlp_dim', default=384, type=int, help='feed forward hidden_dim for a decoder block')
+parser.add_argument('--patch_size', default=16, type=int, help='patch size')
+parser.add_argument('--dim', default=512, type=int, help='embedding dim of patch')
+parser.add_argument('--enc_mlp_dim', default=1024, type=int, help='feed forward hidden_dim for an encoder block')
+parser.add_argument('--dec_mlp_dim', default=1024, type=int, help='feed forward hidden_dim for a decoder block')
 
 
 
@@ -59,8 +59,8 @@ padding_idx = vocab['<PAD>']
 
 model = Transformer(height,width,n_channels,patch_size,dim,enc_head,enc_feed_forward,enc_depth,
                     dec_head,dec_feed_forward,dec_depth,max_seq_len,vocab_size,padding_idx)
-#state = torch.load("captioning.pt",map_location="cpu")
-#model.load_state_dict(state)
+state = torch.load("captioning.ptrom",map_location="cpu")
+model.load_state_dict(state)
 model = model.to(device)
 itos = {v: k for k, v in vocab.items()}