File size: 1,008 Bytes
38c803c 99407d3 38c803c 2a5de8e 20d44a7 38c803c 13e6ee3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
import torch.nn as nn
from transformers import VisionEncoderDecoderModel, PreTrainedModel, AutoConfig
class LegibilityModel(PreTrainedModel):
def __init__(self, config):
config = AutoConfig.from_pretrained("microsoft/trocr-base-handwritten")
super(LegibilityModel, self).__init__(config=config)
# base model architecture
self.model = VisionEncoderDecoderModel(config).encoder
# change dropout during training
self.stack = nn.Sequential(
nn.Dropout(0),
nn.Linear(768, 768),
nn.ReLU(),
nn.Dropout(0),
nn.Linear(768, 1)
)
# choice, img0, img1 are not used by the model, but are passed by the trainer
def forward(self, img_batch, choice=None, img0=None, img1=None):
output = self.model(img_batch)
# average the output of the last hidden layer
output = output.last_hidden_state.mean(dim=1)
scores = self.stack(output)
return scores.squeeze() |