Method readArchiveAndTensors

source/python.js:14339–14379

(archive_name, pickle_prefix, tensor_prefix, type_resolver, obj_loader, device, stream_reader, type_parser, storage_context)

                return this.readArchiveAndTensors(archive_name, this._pickle_dir_prefix, this._tensor_dir_prefix, type_resolver, ObjLoaderFunc, this._device, this._reader, null, this._storage_context);
            }
            readArchiveAndTensors(archive_name, pickle_prefix, tensor_prefix, type_resolver, obj_loader, device, stream_reader, type_parser, storage_context) {
                const picklename = `${pickle_prefix + archive_name}.pkl`;
                const stream = stream_reader.get_record(picklename);
                if (!stream) {
                    throw new python.Error(`File '${picklename}' is not found.`);
                }
                const buffer = stream.peek();
                const tensor_dir_path = tensor_prefix ? tensor_prefix : `${archive_name}/`;
                const read_record = (name) => {
                    const stream = stream_reader.get_record(tensor_dir_path + name);
                    return stream.length <= 0x40000 ? stream.peek() : stream;
                };
                const execution = this._compilation_unit.execution;
                const pickle = execution.__import__('pickle');
                const Unpickler = class extends pickle.Unpickler {
                    find_class(module, name) {
                        return super.find_class(module, name);
                    }
                };
                const unpickler = new Unpickler(buffer);
                unpickler.persistent_load = (saved_id) => {
                    if (saved_id[0] !== 'storage') {
                        throw new python.Error(`Unsupported persistent load type '${saved_id[0]}'.`);
                    }
                    const [, storage_type, key, , size] = saved_id;
                    if (storage_context && storage_context.has_storage(key)) {
                        return storage_context.get_storage(key);
                    }
                    const storage = new storage_type(size);
                    if (!storage._set_cdata) {
                        throw new python.Error(`'${storage_type.__name__}._set_cdata' is not a function.`);
                    }
                    const storage_ptr = read_record(key);
                    storage._set_cdata(storage_ptr);
                    if (storage_context) {
                        storage_context.add_storage(key);
                    }
                    return storage;
                };
                return unpickler.load();
            }
        });
        this.registerType('torch._C.WithInsertPoint', class {
            constructor(...args) {