10
$\begingroup$

I would like to write out the simplest possible word-by-word image caption generator. This network should take an image and build a sentence describing it.

Here is my fashion dataset of images and their string captions. You can download and import it:

Import @ "fashion_captions.wdx"

I need a little help getting the network and training mechanics worked out. Here's my code (and a toy dataset for now) - the definitions of makeDecoderRNN and trainingNet are not right:

vocabSize = 40236; (* from NetEncoder["Tokens"] *)    
embedSize = 50; (* arbitrary for now*)

data = Table[<| 
    "Input" -> RandomImage[1, {224, 224}, ColorSpace -> "RGB"], 
    "Output" -> StringRiffle[RandomWord[5]]|>, 2];

resnet152 = 
  NetTake[NetModel[
    "ResNet-152 Trained on ImageNet Competition Data"], {1, -2}];

makeEncoderCNN[embedSize_] := NetInitialize @ NetChain[
   {resnet152, LinearLayer[embedSize],  BatchNormalizationLayer[]}]

makeDecoderRNN[embedSize_, vocabSize_, lstmSize_: 50, lstmLayers_: 5] := 
  NetInitialize @ NetGraph[<|     
    "lstm" -> NetChain[
       Table[LongShortTermMemoryLayer[lstmSize], lstmLayers]],
    "last" -> SequenceLastLayer[],
    "linear" -> LinearLayer[vocabSize],
    "soft" -> SoftmaxLayer[]
    |>, {
       NetPort["ImageFeatures"] -> NetPort["lstm", {1, "State"}],
       "lstm" -> "last" -> "linear" -> "soft"
    },
   "ImageFeatures" -> embedSize,
   "Output" -> NetDecoder["Tokens"]
   ]

traningNet = NetGraph[
  <|"cnn" -> makeEncoderCNN[embedSize],
   "rnn" -> makeDecoderRNN[embedSize, vocabSize]|>,
  {"cnn" -> NetPort["rnn", "ImageFeatures"]}
  ]

NetTrain[traningNet, data]

Here's a picture of one version of the basic thing I'm trying to implement:

enter image description here

And this is the paper (minus the attention): Show, Attend and Tell: Neural Image Caption Generation with Visual Attention, 2015. I was trying to follow.

The key idea is you somehow inject the image's features into an rnn and iterate one timestep on x0 (start token) trying to output the target y0 (the first word in the caption) then feed that next state and y0 to get y1 (the second word) etc.

Anyhow, I think that there's a NetFoldOperator missing somewhere and may have to use CTC loss?

Links:

Example Data:

$\endgroup$
8
$\begingroup$
enc = NetGraph[
  {
   NetTake[NetModel["ResNet-152 Trained on ImageNet Competition Data"], {1, -2}]
   },
  {
   NetPort["Input"] -> 1 -> NetPort["State"]
   }
  ]

enter image description here

For example, we have 100 tokens + 2 reserved (101 = start of sequence, 102 = end of sequence).

dec = NetGraph[
   {
    EmbeddingLayer[10, 100 + 2],
    SequenceMostLayer[],
    LongShortTermMemoryLayer[1000],
    NetMapOperator[{LinearLayer[100 + 2], SoftmaxLayer[]}]
    },
   {
    NetPort["Input"] -> 1 -> 2 -> 3 -> 4,
    NetPort["State"] -> NetPort[3, "State"]
    }
   ] // NetInitialize

enter image description here

net = NetGraph[
  <|
   "enc" -> enc,
   "dec" -> dec,
   "loss" -> CrossEntropyLossLayer["Index"],
   "rest" -> SequenceRestLayer[]
   |>,
  {
   NetPort["Input"] -> "enc",
   NetPort["enc", "State"] -> NetPort["dec", "State"],
   NetPort["Target"] -> NetPort["dec", "Input"],
   "dec" -> NetPort["loss", "Input"],
   NetPort["Target"] -> "rest" -> NetPort["loss", "Target"]
   },
  "Target" -> {"Varying", "Integer"}
  ]

enter image description here

data = Table[
   RandomImage[1, {224, 224}, ColorSpace -> "RGB"] -> 
    Join[{101}, RandomChoice[Range[100], RandomChoice@Range[10]], {102}],
   {100}
   ];

data[[;; 2]]

enter image description here

predict[input_] := Module[
  {enc, dec, sobj},
  enc = NetReplacePart[NetExtract[netT, "enc"], 
     "Input" -> 
      NetEncoder[{"Image", {224, 224}, ColorSpace -> "RGB"}]]@input;
  dec = RightComposition[
     NetDelete[#, 2] &,
     NetTake[#, {NetPort["Input"], NetPort["Output"]}] &,
     NetReplacePart[#, 
       "Output" -> NetDecoder[{"Class", Range[100 + 2]}]] &
     ]@NetExtract[netT, "dec"];
  sobj = NetStateObject[
    NetTake[dec, {NetPort["Input"], NetPort["Output"]}],
    <|
     {2, "State"} -> enc
     |>
    ];
  Rest@Flatten@NestWhileList[sobj, {101}, # != {102} &, 1, 20]
  ]

Quickly check results with untrained network.

netT = net;

predict /@ data[[;; 2, 1]]

{{24, 24, 24, 24, 13, 15, 15, 50, 50, 20, 20, 20, 8, 8, 8, 8, 8, 8, 8, 8}, {24, 15, 15, 15, 15, 15, 15, 50, 50, 50, 50, 8, 8, 8, 8, 8, 8, 8, 8, 8}}

Train the network.

netT = NetTrain[
  net,
  <|"Input" -> data[[;; , 1]], "Target" -> data[[;; , 2]]|>,
  LearningRateMultipliers -> {"enc" -> None, _ -> 1}
  ]
$\endgroup$
  • $\begingroup$ This is wonderful Alexey! I'm confused on how to use word tokens instead of integers "Target" -> {"Varying", "Token"} doesn't work for me. I added an example dataset to the post. $\endgroup$ – user5601 Sep 21 '18 at 15:20
  • $\begingroup$ I've added a dropbox link to the full dataset (in wdx format) of fashion image -> captions that I'd like to train it on, please take a look if you can. $\endgroup$ – user5601 Sep 21 '18 at 16:08
  • $\begingroup$ With NetEncoder[{"Characters", {Automatic, StartOfString, EndOfString}}] we can encode start and end of sequence. With NetEncoder["Tokens"] we cannot. We need separate encoder and decoder. Encoder: e[str_String]:=Join[(*start of sequence*){40237},NetEncoder[{"Tokens","English","IgnoreCase"->True}]@str,(*end of sequence*){40238}] Decoder: d[list_List] := NetDecoder["Tokens"]@Map[UnitVector[40238, #] &, list] $\endgroup$ – Alexey Golyshev Sep 23 '18 at 14:20
  • $\begingroup$ In dec and predict replace 100 with 40236. The rest of the code is unchanged. data=CloudGet[...] Then predict@data[[1]]["Input"]//d will return text. For training: netT=NetTrain[net,<|"Input"->data[[;;,1]],"Target"->e/@data[[;;,2]]|>,...] $\endgroup$ – Alexey Golyshev Sep 23 '18 at 14:20

Your Answer

By clicking "Post Your Answer", you acknowledge that you have read our updated terms of service, privacy policy and cookie policy, and that your continued use of the website is subject to these policies.

Not the answer you're looking for? Browse other questions tagged or ask your own question.