Как сделать бота, который превращает фото в комикс. Часть третья. Serverless GPU хостинг модели

# Заменяем path на путь к понравившейся итерации модели # Нам нужен файл *G_A.pth - генератор фото -> комикс # output_path - имя файла для экспортируемой модели # может быть любым с расширением *.jit, главное - не забыть, куда мы его сохранили  path= '/checkpoints/resnet9_nowd_nodo_128to400_c8/60_net_G_A.pth' output_path ='/checkpoints/resnet9_nowd_nodo_128to400_c8/resnet9_nowd_nodo_128to400_c8_160-50-60_1.jit'  import torch from torch import nn  class ResnetGenerator(nn.Module):     """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.      We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)     """      def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect'):         """Construct a Resnet-based generator          Parameters:             input_nc (int)      -- the number of channels in input images             output_nc (int)     -- the number of channels in output images             ngf (int)           -- the number of filters in the last conv layer             norm_layer          -- normalization layer             use_dropout (bool)  -- if use dropout layers             n_blocks (int)      -- the number of ResNet blocks             padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero         """         assert(n_blocks >= 0)         super(ResnetGenerator, self).__init__()         if type(norm_layer) == functools.partial:             use_bias = norm_layer.func == nn.InstanceNorm2d         else:             use_bias = norm_layer == nn.InstanceNorm2d          model = [nn.ReflectionPad2d(3),                  nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),                  norm_layer(ngf),                  nn.ReLU(True)]          n_downsampling = 2         for i in range(n_downsampling):  # add downsampling layers             mult = 2 ** i             model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),                       norm_layer(ngf * mult * 2),                       nn.ReLU(True)]          mult = 2 ** n_downsampling         for i in range(n_blocks):       # add ResNet blocks              model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]          for i in range(n_downsampling):  # add upsampling layers             mult = 2 ** (n_downsampling - i)             model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),                                          kernel_size=3, stride=2,                                          padding=1, output_padding=1,                                          bias=use_bias),                       norm_layer(int(ngf * mult / 2)),                       nn.ReLU(True)]         model += [nn.ReflectionPad2d(3)]         model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]         model += [nn.Tanh()]          self.model = nn.Sequential(*model)      def forward(self, input):         """Standard forward"""         return self.model(input)  class ResnetBlock(nn.Module):     """Define a Resnet block"""      def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):         """Initialize the Resnet block          A resnet block is a conv block with skip connections         We construct a conv block with build_conv_block function,         and implement skip connections in <forward> function.         Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf         """         super(ResnetBlock, self).__init__()         self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)      def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):         """Construct a convolutional block.          Parameters:             dim (int)           -- the number of channels in the conv layer.             padding_type (str)  -- the name of padding layer: reflect | replicate | zero             norm_layer          -- normalization layer             use_dropout (bool)  -- if use dropout layers.             use_bias (bool)     -- if the conv layer uses bias or not          Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))         """         conv_block = []         p = 0         if padding_type == 'reflect':             conv_block += [nn.ReflectionPad2d(1)]         elif padding_type == 'replicate':             conv_block += [nn.ReplicationPad2d(1)]         elif padding_type == 'zero':             p = 1         else:             raise NotImplementedError('padding [%s] is not implemented' % padding_type)          conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]         if use_dropout:             conv_block += [nn.Dropout(0.5)]          p = 0         if padding_type == 'reflect':             conv_block += [nn.ReflectionPad2d(1)]         elif padding_type == 'replicate':             conv_block += [nn.ReplicationPad2d(1)]         elif padding_type == 'zero':             p = 1         else:             raise NotImplementedError('padding [%s] is not implemented' % padding_type)         conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]          return nn.Sequential(*conv_block)      def forward(self, x):         """Forward function (with skip connections)"""         out = x + self.conv_block(x)  # add skip connections         return out  import functools  norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)  model = ResnetGenerator(3,3,64, norm_layer=norm_layer, use_dropout=False, n_blocks=9)  model.load_state_dict(torch.load(path)) model.eval() model.cuda() model.half()  trace_input = torch.ones(1,3,256,256).cuda() model = model.eval() jit_model = torch.jit.trace(model.float(), trace_input)  torch.jit.save(jit_model, output_path)

Документация: https://pytorch.org/docs/stable/jit.html
Если быть кратким, экспорт в jit позволяет сериализовать модель и не тащить за собой среду python, все зависимости и внешние модули, которые она могла использовать. Кроме torch, разумеется.
Хоть мы и будем хостить ее в python-среде, jit-модели можно использовать в самостоятельных приложениях.

import Algorithmia  import torch import torchvision import torchvision.transforms as transforms import cv2 from torch import * import uuid import gc  import requests import numpy as np  client = Algorithmia.client()  # Скачиваем модель по ссылке file_path, сохраняем на наш виртуальный сервер # под именем model_file и загружаем ее. def load_model():     file_path = "{ ссылка на нашу загруженную модель }"     model_file = client.file(file_path).getFile().name     model = torch.jit.load(model_file).half().cuda()     return model  model = load_model().eval() torch.cuda.empty_cache() torch.cuda.ipc_collect() torch.backends.cudnn.benchmark = True  # Скачиваем картинку по ссылке, уменьшаем до нужного размера,  # преобразуем в тензор и нормализуем,  # т.к. модель обучалась на нормализованных картинках  def preprocessing(image_path, max_size):      response = requests.get(image_path)     response = response.content     nparr = np.frombuffer(response, np.uint8)     img_res = cv2.imdecode(nparr, cv2.IMREAD_COLOR)     img_res = cv2.cvtColor(img_res, cv2.COLOR_BGR2RGB)      x = img_res.shape[0]     y = img_res.shape[1]      #if image is bigger than the target max_size, downscale it     if x>max_size and x<=y:         y = y*(max_size/x)         x = max_size      if y>max_size and y<x:         x = x*(max_size/y)         y = max_size      size = (int(y),int(x))     img_res = cv2.resize(img_res,size)      t = Tensor(img_res/255.)[None,:]     t = t.permute(0,3,1,2).half().cuda()      # standartize     t = (t-0.5)/0.5      return(t)  def predict(input):      gc.collect()     with torch.no_grad():         res = model(input).detach()      return res  # Денормализуем картинку обратно, сохраняем со случайным именем # и загружаем в наше хранилище   def save_file(res, file_uri):     #de-standartize     res = (res*0.5)+0.5     tempfile = "/tmp/"+str(uuid.uuid4())+".jpg"     torchvision.utils.save_image(res,tempfile)     client.file(file_uri).putFile(tempfile)  # API calls will begin at the apply() method, with the request body passed as 'input' # For more details, see algorithmia.com/developers/algorithm-development/languages def apply(input):      processed_data= preprocessing(input["in"], input["size"])     res = predict(processed_data)     save_file(res, input["out"])     input = None     res = None     processed_data = None      gc.collect()      return "Success"