Wie man einen Bot macht, der ein Foto in einen Comic verwandelt. Teil drei. Kostenloses serverloses + GPU-Hosting-Modell

#  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) 

Dokumentation: https://pytorch.org/docs/stable/jit.html
Kurz gesagt, beim Exportieren nach JIT können Sie das Modell serialisieren und nicht die Python-Umgebung, alle Abhängigkeiten und externen Module, die verwendet werden könnten, mitziehen. Außer Fackel natürlich.
Obwohl wir es in einer Python-Umgebung hosten werden, können JIT-Modelle in eigenständigen Anwendungen verwendet werden.

 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"