Cómo hacer un bot que convierta una foto en un cómic. Tercera parte Modelo de alojamiento gratuito sin servidor + 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) 

Documentación: https://pytorch.org/docs/stable/jit.html
En resumen, exportar a jit le permite serializar el modelo y no arrastrar el entorno de Python junto con él, todas las dependencias y módulos externos que podría usar. Excepto la antorcha, por supuesto.
Aunque lo alojaremos en un entorno python, los modelos jit se pueden usar en aplicaciones independientes.

 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"